[autotest] Support for new autotest artifacts.

This change updates dynamic_suites and provisioning to support the new
autotest build artifacts: control_files.tar and autotest_packages.tar.

In case these artifacts do not exist, it will then attempt to stage the
old autotest.tar artifact, thus still supporting older builds.

BUG=chromium:421122
TEST=Custom moblab build that successfully ran smoke suite against
2 images: 1 with autotest.tar and the other with the new artifacts.
CQ-DEPEND=CL:226822

Change-Id: If66026b46f77aaccaa83996e20a6749bad2b1625
Reviewed-on: https://chromium-review.googlesource.com/226823
Reviewed-by: Simran Basi <sbasi@chromium.org>
Commit-Queue: Simran Basi <sbasi@chromium.org>
Tested-by: Simran Basi <sbasi@chromium.org>
diff --git a/client/common_lib/cros/dev_server.py b/client/common_lib/cros/dev_server.py
index 3e2c978..8e08354 100644
--- a/client/common_lib/cros/dev_server.py
+++ b/client/common_lib/cros/dev_server.py
@@ -40,7 +40,12 @@
 # Artifacts that should be staged when client calls devserver RPC to stage an
 # image with autotest artifact.
 _ARTIFACTS_TO_BE_STAGED_FOR_IMAGE_WITH_AUTOTEST = ('full_payload,test_suites,'
-                                                   'autotest,stateful')
+                                                   'control_files,stateful,'
+                                                   'autotest_packages')
+# This dictionary translates newer smaller artifacts to their old compatible
+# artifact. This allows us to continue to support old builds.
+_COMPATIBLE_ARTIFACTS = {'control_files': 'autotest',
+                         'autotest_packages' : 'autotest'}
 
 
 class MarkupStripper(HTMLParser.HTMLParser):
@@ -130,6 +135,50 @@
     return inner_decorator
 
 
+# TODO (sbasi) crbug.com/433436 - Remove this decorator once we no longer care
+# about old builds that have the old autotest artifact.
+def compatible_artifacts():
+    """A decorator to use with the stage_artifact call.
+
+    Over time some artifacts may be replaced with other compatible artifacts.
+    This decorator retries calls involving such artifacts with the
+    compatible artifact instead.
+    """
+    #pylint: disable=C0111
+    def inner_decorator(method):
+        def wrapper(ds, image, artifacts=None, files=None, archive_url=None):
+            """
+            @param ds: DevServer instance.
+            @param image: image we require artifacts for.
+            @param artifacts: List of artifacts we want to stage.
+            @param **kwargs: remaining args passed to decorated method.
+            @param files: A list of files to stage.
+            @param archive_url: Optional parameter that has the archive_url to
+                   stage this artifact from. Default is specified in autotest
+                   config + image.
+
+            @raise DevServerException should there not be comparable artifacts.
+            """
+            try:
+                method(ds, image, artifacts=artifacts, files=files,
+                       archive_url=archive_url)
+            except DevServerException as e:
+                if not artifacts or not set.intersection(
+                        set(artifacts), set(_COMPATIBLE_ARTIFACTS.keys())):
+                    raise e
+                logging.debug('Failed to stage %s for %s: %s', artifacts,
+                              image, e)
+                artifacts = [_COMPATIBLE_ARTIFACTS.get(x,x) for x in artifacts]
+                logging.debug('Trying to stage compatible artifacts: %s',
+                              artifacts)
+                method(ds, image, artifacts=artifacts, files=files,
+                       archive_url=archive_url)
+
+        return wrapper
+
+    return inner_decorator
+
+
 class DevServerException(Exception):
     """Raised when the dev server returns a non-200 HTTP response."""
     pass
@@ -521,12 +570,13 @@
         return metadata
 
 
+    @compatible_artifacts()
     @remote_devserver_call()
     def stage_artifacts(self, image, artifacts=None, files=None,
                         archive_url=None):
         """Tell the devserver to download and stage |artifacts| from |image|.
 
-        This is the main call point for staging any specific artifacts for a
+         This is the main call point for staging any specific artifacts for a
         given build. To see the list of artifacts one can stage see:
 
         ~src/platfrom/dev/artifact_info.py.
diff --git a/contrib/stage_build.py b/contrib/stage_build.py
index ea9fd0f..d433210 100755
--- a/contrib/stage_build.py
+++ b/contrib/stage_build.py
@@ -37,8 +37,8 @@
     ds = dev_server.ImageServer.resolve(options.build)
 
   print "Downloading %s..." % options.build
-  ds.stage_artifacts(options.build, ['full_payload', 'stateful', 'autotest'])
-
+  ds.stage_artifacts(options.build, ['full_payload', 'stateful',
+                                     'control_files', 'autotest_packages'])
   if options.host:
     print "Poking job_repo_url on %s..." % options.host
     repo_url = tools.get_package_url(ds.url(), options.build)
diff --git a/server/cros/dynamic_suite/dynamic_suite.py b/server/cros/dynamic_suite/dynamic_suite.py
index 0128105..aa61586 100644
--- a/server/cros/dynamic_suite/dynamic_suite.py
+++ b/server/cros/dynamic_suite/dynamic_suite.py
@@ -464,7 +464,8 @@
     # We can't do anything else until the devserver has finished downloading
     # autotest.tar so that we can get the control files we should schedule.
     try:
-        spec.devserver.stage_artifacts(spec.build, ['autotest', 'test_suites'])
+        spec.devserver.stage_artifacts(
+                spec.build, ['control_files', 'test_suites'])
     except dev_server.DevServerException as e:
         # If we can't get the control files, there's nothing to run.
         raise error.AsynchronousBuildFailure(e)
diff --git a/server/cros/dynamic_suite/dynamic_suite_unittest.py b/server/cros/dynamic_suite/dynamic_suite_unittest.py
index a1cd48f..2ccaa6a 100755
--- a/server/cros/dynamic_suite/dynamic_suite_unittest.py
+++ b/server/cros/dynamic_suite/dynamic_suite_unittest.py
@@ -138,7 +138,7 @@
         spec.build = ''
         spec.devserver = self.mox.CreateMock(dev_server.ImageServer)
         spec.devserver.stage_artifacts(
-                spec.build, ['autotest', 'test_suites']).WithSideEffects(
+                spec.build, ['control_files', 'test_suites']).WithSideEffects(
                 suicide)
 
         self.mox.ReplayAll()
diff --git a/server/hosts/cros_host.py b/server/hosts/cros_host.py
index 63eac75..fcaedaf 100644
--- a/server/hosts/cros_host.py
+++ b/server/hosts/cros_host.py
@@ -456,7 +456,7 @@
             image_name, ds.url())
 
         start_time = time.time()
-        ds.stage_artifacts(image_name, ['autotest'])
+        ds.stage_artifacts(image_name, ['autotest_packages'])
         stage_time = time.time() - start_time
 
         # Record how much of the verification time comes from a devserver
diff --git a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
index 8724338..9447423 100644
--- a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
+++ b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
@@ -59,7 +59,8 @@
         # reimaging finishes or at some other point in the provisioning.
         try:
             ds = dev_server.ImageServer.resolve(image)
-            ds.stage_artifacts(image, ['full_payload', 'stateful', 'autotest'])
+            ds.stage_artifacts(image, ['full_payload', 'stateful',
+                                       'autotest_packages'])
         except dev_server.DevServerException as e:
             raise error.TestFail(str(e))