Shows the image when clone a suite job or its child job.

BUG=chromium:485813
TEST=unit test and manual test.

Change-Id: I044991928513cea2c06eb6db8b31eeba6ea29289
Reviewed-on: https://chromium-review.googlesource.com/348482
Commit-Ready: Michael Tang <ntang@chromium.org>
Tested-by: Michael Tang <ntang@chromium.org>
Reviewed-by: Simran Basi <sbasi@chromium.org>
Reviewed-by: Michael Tang <ntang@chromium.org>
diff --git a/client/common_lib/control_data.py b/client/common_lib/control_data.py
index 76c5e65..0463354 100644
--- a/client/common_lib/control_data.py
+++ b/client/common_lib/control_data.py
@@ -258,6 +258,14 @@
         self._set_bool('require_ssp', val)
 
 
+    def set_build(self, val):
+        self._set_string('build', val)
+
+
+    def set_builds(self, val):
+        if type(val) == dict:
+            setattr(self, 'builds', val)
+
     def set_attributes(self, val):
         # Add subsystem:default if subsystem is not specified.
         self._set_set('attributes', val)
diff --git a/frontend/afe/rpc_interface.py b/frontend/afe/rpc_interface.py
index f7c8dcb..dff4165 100644
--- a/frontend/afe/rpc_interface.py
+++ b/frontend/afe/rpc_interface.py
@@ -31,12 +31,14 @@
 
 __author__ = 'showard@google.com (Steve Howard)'
 
+import ast
 import sys
 import datetime
 import logging
 
 from django.db.models import Count
 import common
+from autotest_lib.client.common_lib import control_data
 from autotest_lib.client.common_lib import priorities
 from autotest_lib.client.common_lib.cros import dev_server
 from autotest_lib.client.common_lib.cros.graphite import autotest_stats
@@ -890,7 +892,7 @@
 def create_job_page_handler(name, priority, control_file, control_type,
                             image=None, hostless=False, firmware_rw_build=None,
                             firmware_ro_build=None, test_source_build=None,
-                            **kwargs):
+                            is_cloning = False, **kwargs):
     """\
     Create and enqueue a job.
 
@@ -905,11 +907,16 @@
                               None, i.e., RO firmware will not be updated.
     @param test_source_build: Build to be used to retrieve test code. Default
                               to None.
+    @param is_cloning: True if creating a cloning job.
     @param kwargs extra args that will be required by create_suite_job or
                   create_job.
 
     @returns The created Job id number.
     """
+    if is_cloning:
+        logging.info('Start to clone a new job')
+    else:
+        logging.info('Start to create a new job')
     control_file = rpc_utils.encode_ascii(control_file)
     if not control_file:
         raise model_logic.ValidationError({
@@ -924,9 +931,10 @@
             builds[provision.FW_RO_VERSION_PREFIX] = firmware_ro_build
         return site_rpc_interface.create_suite_job(
                 name=name, control_file=control_file, priority=priority,
-                builds=builds, test_source_build=test_source_build, **kwargs)
+                builds=builds, test_source_build=test_source_build,
+                is_cloning=is_cloning, **kwargs)
     return create_job(name, priority, control_file, control_type, image=image,
-                      hostless=hostless, **kwargs)
+                      hostless=hostless, is_cloning=is_cloning, **kwargs)
 
 
 @rpc_utils.route_rpc_to_master
@@ -938,7 +946,7 @@
                reboot_before=None, reboot_after=None, parse_failed_repair=None,
                hostless=False, keyvals=None, drone_set=None, image=None,
                parent_job_id=None, test_retry=0, run_reset=True,
-               require_ssp=None, args=(), **kwargs):
+               require_ssp=None, args=(), is_cloning=False, **kwargs):
     """\
     Create and enqueue a job.
 
@@ -981,6 +989,7 @@
                        image is not set, drone will run the test without server-
                        side packaging. Default is None.
     @param args A list of args to be injected into control file.
+    @param is_cloning: True if creating a cloning job.
     @param kwargs extra keyword args. NOT USED.
 
     @returns The created Job id number.
@@ -1266,13 +1275,58 @@
     info['hostless'] = job_info['hostless']
     info['drone_set'] = job.drone_set and job.drone_set.name
 
-    if job.parameterized_job:
-        info['job']['image'] = get_parameterized_autoupdate_image_url(job)
+    image = _get_image_for_job(job, job_info['hostless'])
+    if image:
+        info['job']['image'] = image
 
     return rpc_utils.prepare_for_serialization(info)
 
 
-# host queue entries
+def _get_image_for_job(job, hostless):
+    """ Gets the image used for a job.
+
+    Gets the image used for an AFE job. If the job is a parameterized job, get
+    the image from the job parameter; otherwise, tries to get the image from
+    the job's keyvals 'build' or 'builds'. As a last resort, if the job is a
+    hostless job, tries to get the image from its control file attributes
+    'build' or 'builds'.
+
+    TODO(ntang): Needs to handle FAFT with two builds for ro/rw.
+
+    @param job      An AFE job object.
+    @param hostless Boolean on of the job is hostless.
+
+    @returns The image build used for the job.
+    """
+    image = None
+    if job.parameterized_job:
+        image = get_parameterized_autoupdate_image_url(job)
+    else:
+        keyvals = job.keyval_dict()
+        image =  keyvals.get('build')
+        if not image:
+            value = keyvals.get('builds')
+            builds = None
+            if isinstance(value, dict):
+                builds = value
+            elif isinstance(value, basestring):
+                builds = ast.literal_eval(value)
+            if builds:
+                image = builds.get('cros-version')
+        if not image and hostless and job.control_file:
+            try:
+                control_obj = control_data.parse_control_string(
+                        job.control_file)
+                if hasattr(control_obj, 'build'):
+                    image = getattr(control_obj, 'build')
+                if not image and hasattr(control_obj, 'builds'):
+                    builds = getattr(control_obj, 'builds')
+                    image = builds.get('cros-version')
+            except:
+                logging.warning('Failed to parse control file for job: %s',
+                                job.name)
+    return image
+
 
 def get_host_queue_entries(start_time=None, end_time=None, **filter_data):
     """\
diff --git a/frontend/afe/rpc_interface_unittest.py b/frontend/afe/rpc_interface_unittest.py
index 0cb5cb0..7a73860 100755
--- a/frontend/afe/rpc_interface_unittest.py
+++ b/frontend/afe/rpc_interface_unittest.py
@@ -629,5 +629,74 @@
         self.god.check_playback()
 
 
+    def test_get_image_for_job_parameterized(self):
+        test = models.Test.objects.create(
+            name='name', author='author', test_class='class',
+            test_category='category',
+            test_type=control_data.CONTROL_TYPE.SERVER, path='path')
+        parameterized_job = models.ParameterizedJob.objects.create(test=test)
+        job = self._create_job(hosts=[1])
+        job.parameterized_job = parameterized_job
+        self.god.stub_function_to_return(rpc_interface,
+                'get_parameterized_autoupdate_image_url', 'cool-image')
+        image = rpc_interface._get_image_for_job(job, True)
+        self.assertEquals('cool-image', image)
+        self.god.check_playback()
+
+
+    def test_get_image_for_job_with_keyval_build(self):
+        keyval_dict = {'build': 'cool-image'}
+        job_id = rpc_interface.create_job(name='test', priority='Medium',
+                                          control_file='foo',
+                                          control_type=CLIENT,
+                                          hosts=['host1'],
+                                          keyvals=keyval_dict)
+        job = models.Job.objects.get(id=job_id)
+        self.assertIsNotNone(job)
+        image = rpc_interface._get_image_for_job(job, True)
+        self.assertEquals('cool-image', image)
+
+
+    def test_get_image_for_job_with_keyval_builds(self):
+        keyval_dict = {'builds': {'cros-version': 'cool-image'}}
+        job_id = rpc_interface.create_job(name='test', priority='Medium',
+                                          control_file='foo',
+                                          control_type=CLIENT,
+                                          hosts=['host1'],
+                                          keyvals=keyval_dict)
+        job = models.Job.objects.get(id=job_id)
+        self.assertIsNotNone(job)
+        image = rpc_interface._get_image_for_job(job, True)
+        self.assertEquals('cool-image', image)
+
+
+    def test_get_image_for_job_with_control_build(self):
+        CONTROL_FILE = """build='cool-image'
+        """
+        job_id = rpc_interface.create_job(name='test', priority='Medium',
+                                          control_file='foo',
+                                          control_type=CLIENT,
+                                          hosts=['host1'])
+        job = models.Job.objects.get(id=job_id)
+        self.assertIsNotNone(job)
+        job.control_file = CONTROL_FILE
+        image = rpc_interface._get_image_for_job(job, True)
+        self.assertEquals('cool-image', image)
+
+
+    def test_get_image_for_job_with_control_builds(self):
+        CONTROL_FILE = """builds={'cros-version': 'cool-image'}
+        """
+        job_id = rpc_interface.create_job(name='test', priority='Medium',
+                                          control_file='foo',
+                                          control_type=CLIENT,
+                                          hosts=['host1'])
+        job = models.Job.objects.get(id=job_id)
+        self.assertIsNotNone(job)
+        job.control_file = CONTROL_FILE
+        image = rpc_interface._get_image_for_job(job, True)
+        self.assertEquals('cool-image', image)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/frontend/afe/site_rpc_interface.py b/frontend/afe/site_rpc_interface.py
index a629f58..6077326 100644
--- a/frontend/afe/site_rpc_interface.py
+++ b/frontend/afe/site_rpc_interface.py
@@ -153,7 +153,7 @@
                      max_retries=None, max_runtime_mins=None, suite_min_duts=0,
                      offload_failures_only=False, builds={},
                      test_source_build=None, run_prod_code=False,
-                     delay_minutes=0, **kwargs):
+                     delay_minutes=0, is_cloning=False, **kwargs):
     """
     Create a job to run a test suite on the given device with the given image.
 
@@ -201,6 +201,7 @@
                           build artifacts.
     @param delay_minutes: Delay the creation of test jobs for a given number of
                           minutes.
+    @param is_cloning: True if creating a cloning job.
     @param kwargs: extra keyword args. NOT USED.
 
     @raises ControlFileNotFound: if a unique suite control file doesn't exist.
@@ -292,6 +293,8 @@
                    'delay_minutes': delay_minutes,
                    }
 
+    if is_cloning:
+        control_file = tools.remove_injection(control_file)
     control_file = tools.inject_vars(inject_dict, control_file)
 
     return rpc_utils.create_job_common(name,
@@ -1029,7 +1032,7 @@
         try:
             control_obj = control_data.parse_control_string(control_file)
         except:
-            logging.info('Failed to parse congtrol file: %s', control_file_path)
+            logging.info('Failed to parse control file: %s', control_file_path)
             if not ignore_invalid_tests:
                 raise
 
diff --git a/frontend/client/src/autotest/afe/AfeClient.java b/frontend/client/src/autotest/afe/AfeClient.java
index 2380396..257810c 100644
--- a/frontend/client/src/autotest/afe/AfeClient.java
+++ b/frontend/client/src/autotest/afe/AfeClient.java
@@ -86,8 +86,10 @@
 
             public void onCloneJob(JSONValue cloneInfo) {
                 createJob.ensureInitialized();
-                createJob.cloneJob(cloneInfo);
                 mainTabPanel.selectTabView(createJob);
+                // Makes sure we set cloning mode after the tab is selected. Otherwise,
+                // the mode will be reset.
+                createJob.cloneJob(cloneInfo);
             }
 
             public void onCreateRecurringJob(int jobId) {
diff --git a/frontend/client/src/autotest/afe/create/CreateJobViewPresenter.java b/frontend/client/src/autotest/afe/create/CreateJobViewPresenter.java
index 1f149a4..0a37005 100644
--- a/frontend/client/src/autotest/afe/create/CreateJobViewPresenter.java
+++ b/frontend/client/src/autotest/afe/create/CreateJobViewPresenter.java
@@ -149,6 +149,7 @@
     private JSONArray dependencies = new JSONArray();
 
     private Display display;
+    private boolean cloning;
 
     public void bindDisplay(Display display) {
         this.display = display;
@@ -158,7 +159,16 @@
         this.listener = listener;
     }
 
+    public void setCloning(boolean cloning) {
+      this.cloning = cloning;
+    }
+
+    public boolean isCloning() {
+      return cloning;
+    }
+
     public void cloneJob(JSONValue cloneInfo) {
+        setCloning(true);
         // reset() fires the TestSelectorListener, which will generate a new control file. We do
         // no want this, so we'll stop listening to it for a bit.
         testSelector.setListener(null);
@@ -825,6 +835,7 @@
                     args.put(TEST_SOURCE_BUILD, new JSONString(testSourceBuild));
                 }
 
+                args.put("is_cloning", JSONBoolean.getInstance(isCloning()));
                 rpcProxy.rpcCall("create_job_page_handler", args, new JsonRpcCallback() {
                     @Override
                     public void onSuccess(JSONValue result) {
diff --git a/frontend/client/src/autotest/afe/create/CreateJobViewTab.java b/frontend/client/src/autotest/afe/create/CreateJobViewTab.java
index 8a30d40..fd2701a 100644
--- a/frontend/client/src/autotest/afe/create/CreateJobViewTab.java
+++ b/frontend/client/src/autotest/afe/create/CreateJobViewTab.java
@@ -24,6 +24,13 @@
     }
 
     @Override
+    public void ensureInitialized() {
+      super.ensureInitialized();
+      // Makes sure cloning mode is turned off.
+      getPresenter().setCloning(false);
+    }
+
+    @Override
     public void initialize() {
         super.initialize();
         getDisplay().initialize((HTMLPanel) getWidget());
diff --git a/server/cros/dynamic_suite/tools.py b/server/cros/dynamic_suite/tools.py
index 11507ae..9db611f 100644
--- a/server/cros/dynamic_suite/tools.py
+++ b/server/cros/dynamic_suite/tools.py
@@ -14,6 +14,16 @@
 _CONFIG = global_config.global_config
 
 
+# comments injected into the control file.
+_INJECT_BEGIN = '# INJECT_BEGIN - DO NOT DELETE THIS LINE'
+_INJECT_END = '# INJECT_END - DO NOT DELETE LINE'
+
+
+# The regex for an injected line in the control file with the format:
+# varable_name=varable_value
+_INJECT_VAR_RE = re.compile('^[_A-Za-z]\w*=.+$')
+
+
 def image_url_pattern():
     """Returns image_url_pattern from global_config."""
     return _CONFIG.get_config_value('CROS', 'image_url_pattern', type=str)
@@ -144,6 +154,58 @@
     return None
 
 
+def remove_legacy_injection(control_file_in):
+    """
+    Removes the legacy injection part from a control file.
+
+    @param control_file_in: the contents of a control file to munge.
+
+    @return The modified control file string.
+    """
+    if not control_file_in:
+        return control_file_in
+
+    new_lines = []
+    lines = control_file_in.strip().splitlines()
+    remove_done = False
+    for line in lines:
+        if remove_done:
+            new_lines.append(line)
+        else:
+            if not _INJECT_VAR_RE.match(line):
+                remove_done = True
+                new_lines.append(line)
+    return '\n'.join(new_lines)
+
+
+def remove_injection(control_file_in):
+    """
+    Removes the injection part from a control file.
+
+    @param control_file_in: the contents of a control file to munge.
+
+    @return The modified control file string.
+    """
+    if not control_file_in:
+        return control_file_in
+
+    start = control_file_in.find(_INJECT_BEGIN)
+    if start >=0:
+        end = control_file_in.find(_INJECT_END, start)
+    if start < 0 or end < 0:
+        return remove_legacy_injection(control_file_in)
+
+    end += len(_INJECT_END)
+    ch = control_file_in[end]
+    total_length = len(control_file_in)
+    while end <= total_length and (
+            ch == '\n' or ch == ' ' or ch == '\t'):
+        end += 1
+        if end < total_length:
+          ch = control_file_in[end]
+    return control_file_in[:start] + control_file_in[end:]
+
+
 def inject_vars(vars, control_file_in):
     """
     Inject the contents of |vars| into |control_file_in|.
@@ -153,6 +215,7 @@
     @return the modified control file string.
     """
     control_file = ''
+    control_file += _INJECT_BEGIN + '\n'
     for key, value in vars.iteritems():
         # None gets injected as 'None' without this check; same for digits.
         if isinstance(value, str):
@@ -161,7 +224,7 @@
             control_file += "%s=%r\n" % (key, value)
 
     args_dict_str = "%s=%s\n" % ('args_dict', repr(vars))
-    return control_file + args_dict_str + control_file_in
+    return control_file + args_dict_str + _INJECT_END + '\n' + control_file_in
 
 
 def is_usable(host):
diff --git a/server/cros/dynamic_suite/tools_unittest.py b/server/cros/dynamic_suite/tools_unittest.py
index 89d29c6..dcedda5 100755
--- a/server/cros/dynamic_suite/tools_unittest.py
+++ b/server/cros/dynamic_suite/tools_unittest.py
@@ -42,7 +42,11 @@
         """Should inject dict of varibles into provided strings."""
         def find_all_in(d, s):
             """Returns true if all key-value pairs in |d| are printed in |s|
-            and the dictionary representation is also in |s|."""
+            and the dictionary representation is also in |s|.
+
+            @param d: the variable dictionary to check.
+            @param s: the control file string.
+            """
             for k, v in d.iteritems():
                 if isinstance(v, str):
                     if "%s='%s'\n" % (k, v) not in s:
@@ -58,6 +62,39 @@
         v = {'v1': 'one', 'v2': 'two', 'v3': None, 'v4': False, 'v5': 5}
         self.assertTrue(find_all_in(v, tools.inject_vars(v, '')))
         self.assertTrue(find_all_in(v, tools.inject_vars(v, 'ctrl')))
+        control_file = tools.inject_vars(v, 'sample')
+        self.assertTrue(tools._INJECT_BEGIN in control_file)
+        self.assertTrue(tools._INJECT_END in control_file)
+
+    def testRemoveInjection(self):
+        """Tests remove the injected variables from control file."""
+        control_file = """
+# INJECT_BEGIN - DO NOT DELETE THIS LINE
+v1='one'
+v4=False
+v5=5
+args_dict={'v1': 'one', 'v2': 'two', 'v3': None, 'v4': False, 'v5': 5}
+# INJECT_END - DO NOT DELETE LINE
+def init():
+    pass
+        """
+        control_file = tools.remove_injection(control_file)
+        self.assertTrue(control_file.strip().startswith('def init():'))
+
+    def testRemoveLegacyInjection(self):
+        """Tests remove the injected variables from legacy control file."""
+        control_file = """
+v1='one'
+_v2=False
+v3_x11_=5
+args_dict={'v1': 'one', '_v2': False, 'v3_x11': 5}
+def init():
+    pass
+        """
+        control_file = tools.remove_legacy_injection(control_file)
+        self.assertTrue(control_file.strip().startswith('def init():'))
+        control_file = tools.remove_injection(control_file)
+        self.assertTrue(control_file.strip().startswith('def init():'))
 
 
     def testIncorrectlyLocked(self):