[Android] Add sharding for AMP instrumentation tests.

build/android/pylib/remote/device/remote_device_instrumentation_test_run.py

Index: build/android/pylib/remote/device/remote_device_instrumentation_test_run.py
diff --git a/build/android/pylib/remote/device/remote_device_instrumentation_test_run.py b/build/android/pylib/remote/device/remote_device_instrumentation_test_run.py
index 8858b069ce6fdd4bed80cdbef419f7a7e3f32e06..0326c3aac266430f5cb82a14719bed5f15597db4 100644
--- a/build/android/pylib/remote/device/remote_device_instrumentation_test_run.py
+++ b/build/android/pylib/remote/device/remote_device_instrumentation_test_run.py
@@ -7,7 +7,11 @@
 import logging
 import os
 import tempfile
+import zipfile
 
+from devil.utils import reraiser_thread
+from devil.utils import tempdir
+from devil.utils import zip_utils
 from pylib.base import base_test_result
 from pylib.remote.device import remote_device_test_run
 
@@ -21,55 +25,108 @@ class RemoteDeviceInstrumentationTestRun(
     return self._test_instance.test_package
 
   #override
-  def _TriggerSetUp(self):
-    """Set up the triggering of a test run."""
-    logging.info('Triggering test run.')
-
-    # pylint: disable=protected-access
-    with tempfile.NamedTemporaryFile(suffix='.txt') as test_list_file:
-      tests = self._test_instance.GetTests()
-      logging.debug('preparing to run %d instrumentation tests remotely:',
-                    len(tests))
-      for t in tests:
-        test_name = '%s#%s' % (t['class'], t['method'])
-        logging.debug('  %s', test_name)
-        test_list_file.write('%s\n' % test_name)
-      test_list_file.flush()
-      self._test_instance._data_deps.append(
-          (os.path.abspath(test_list_file.name), None))
-
-      env_vars = self._test_instance.GetDriverEnvironmentVars(
-          test_list_file_path=test_list_file.name)
-      env_vars.update(self._test_instance.GetHttpServerEnvironmentVars())
-
-      logging.debug('extras:')
-      for k, v in env_vars.iteritems():
-        logging.debug('  %s: %s', k, v)
-
-      self._AmInstrumentTestSetup(
-          self._test_instance.apk_under_test,
-          self._test_instance.driver_apk,
-          self._test_instance.driver_name,
-          environment_variables=env_vars,
-          extra_apks=([self._test_instance.test_apk] +
-                      self._test_instance.additional_apks))
+  def _GetAppPath(self):
+    return self._test_instance.apk_under_test
 
   #override
-  def _ParseTestResults(self):
+  def _GetTestFramework(self):
+    if self._env.test_framework:
+      logging.warning('Ignoring configured test_framework "%s"',
+                      self._env.test_framework)
+    return 'robotium'
+
+  #override
+  def _ParseTestResults(self, test_output, results_zip):
     logging.info('Parsing results from stdout.')
-    r = base_test_result.TestRunResults()
+    results = base_test_result.TestRunResults()
     result_code, result_bundle, statuses = (
         self._test_instance.ParseAmInstrumentRawOutput(
-            self._results['results']['output'].splitlines()))
+            test_output['results']['output'].splitlines()))
     result = self._test_instance.GenerateTestResults(
         result_code, result_bundle, statuses, 0, 0)
 
     if isinstance(result, base_test_result.BaseTestResult):
-      r.AddResult(result)
+      results.AddResult(result)
     elif isinstance(result, list):
-      r.AddResults(result)
+      results.AddResults(result)
     else:
       raise Exception('Unexpected result type: %s' % type(result).__name__)
 
-    self._DetectPlatformErrors(r)
-    return r
+    remote_device_test_run.DetectPlatformErrors(
+        results, test_output, results_zip)
+    return results
+
+  #override
+  def _ShouldShard(self):
+    return True
+
+  #override
+  def _SetupTestShards(self, num_shards):
+    test_ids = []
+
+    all_tests = self._test_instance.GetTests()
+    test_shards = [all_tests[i::num_shards] for i in range(num_shards)]
+
+    test_list_files = []
+    with tempdir.TempDir() as test_list_file_dir:

[rnephew (Wrong account), 2015/11/04 22:22:03]

Just a suggestion, feel free to ignore:
This seems a bit verbose. Maybe just temp_dir, or test_temp_dir?

[mikecase (-- gone --), 2015/11/05 23:17:41]

Changed to test_list_dir.

+      for i, tests in enumerate(test_shards):
+        test_list_file = tempfile.NamedTemporaryFile(
+            suffix='.txt', dir=test_list_file_dir, delete=False)
+        test_list_files.append(test_list_file)
+        for t in tests:
+          test_name = '%s#%s' % (t['class'], t['method'])
+          logging.debug('shard %s: %s', i, test_name)
+          test_list_file.write('%s\n' % test_name)
+        test_list_file.flush()
+
+      with tempfile.NamedTemporaryFile(suffix='.zip') as test_package:
+        extra_apks = ([self._test_instance.test_apk] +
+                      self._test_instance.additional_apks)
+        data_deps = (self._test_instance.GetDataDependencies() +
+                     [(os.path.abspath(f.name), None) for f in test_list_files])
+
+        sdcard_files = []
+        additional_apks = []
+        host_test = os.path.basename(self._test_instance.driver_apk)
+        with zipfile.ZipFile(test_package.name, 'w') as zip_file:
+          zip_file.write(
+              self._test_instance.driver_apk, host_test, zipfile.ZIP_DEFLATED)
+          for h, _ in data_deps:
+            if os.path.isdir(h):
+              zip_utils.WriteToZipFile(zip_file, h, '.')
+              sdcard_files.extend(os.listdir(h))
+            else:
+              zip_utils.WriteToZipFile(zip_file, h, os.path.basename(h))
+              sdcard_files.append(os.path.basename(h))
+          for a in extra_apks:
+            zip_utils.WriteToZipFile(zip_file, a, os.path.basename(a))
+            additional_apks.append(os.path.basename(a))
+
+        framework_configs = {
+          'runner': self._test_instance.driver_name,
+          'sdcard_files': ','.join(sdcard_files),
+          'host_test': host_test,
+          'additional_apks': ','.join(additional_apks),
+        }
+
+        def upload_test_shards(test_list_file):
+          shard_framework_config = dict(framework_configs)
+          env_vars = self._test_instance.GetDriverEnvironmentVars(
+              test_list_file_path=test_list_file.name)
+          env_vars.update(self._test_instance.GetHttpServerEnvironmentVars())
+          shard_framework_config['environment_vars'] = ','.join(
+              '%s=%s' % (k, v) for k, v in env_vars.iteritems())
+
+          test_id = self._UploadTestToDevice(test_package.name)
+          self._UploadTestConfigToDevice(
+              test_id, shard_framework_config, self._appurify_configs)
+          test_ids.append(test_id)
+
+        workers = reraiser_thread.ReraiserThreadGroup(
+            [reraiser_thread.ReraiserThread(
+                upload_test_shards, args=[test_list_file],
+                name='upload test shard %d' % i)
+             for i, test_list_file in enumerate(test_list_files)])
+        workers.StartAll()
+        workers.JoinAll()
+    return test_ids