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/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):