Retreive screenshots captured by tests from test devices.

build/android/pylib/local/device/local_device_instrumentation_test_run.py

Index: build/android/pylib/local/device/local_device_instrumentation_test_run.py
diff --git a/build/android/pylib/local/device/local_device_instrumentation_test_run.py b/build/android/pylib/local/device/local_device_instrumentation_test_run.py
index da4da4027a1e16772665cc8ae55f3ae485aecdf1..906487f009dc5892fb86554ee124a903ccce301f 100644
--- a/build/android/pylib/local/device/local_device_instrumentation_test_run.py
+++ b/build/android/pylib/local/device/local_device_instrumentation_test_run.py
@@ -55,6 +55,11 @@ LOGCAT_FILTERS = ['*:e', 'chromium:v', 'cr_*:v', 'DEBUG:I']
 EXTRA_SCREENSHOT_FILE = (
     'org.chromium.base.test.ScreenshotOnFailureStatement.ScreenshotFile')
 
+EXTRA_UI_CAPTURE_DIR = (
+    'org.chromium.base.test.util.Screenshooter.ScreenshotDir')
+
+UI_CAPTURE_DIRS = ['chromium_tests_root', 'UiCapture']
+
 FEATURE_ANNOTATION = 'Feature'
 RENDER_TEST_FEATURE_ANNOTATION = 'RenderTest'
 
@@ -92,6 +97,7 @@ class LocalDeviceInstrumentationTestRun(
   def __init__(self, env, test_instance):
     super(LocalDeviceInstrumentationTestRun, self).__init__(env, test_instance)
     self._flag_changers = {}
+    self._ui_capture_dir = dict()
 
   #override
   def TestPackage(self):
@@ -216,8 +222,20 @@ class LocalDeviceInstrumentationTestRun(
         valgrind_tools.SetChromeTimeoutScale(
             dev, self._test_instance.timeout_scale)
 
+      @trace_event.traced
+      def setup_ui_capture_dir():
+        # Make sure the UI capture directory exists and is empty by deleting
+        # and recreating it.
+        # TODO (aberent) once DeviceTempDir exists use it here.
+        self._ui_capture_dir[dev] = posixpath.join(dev.GetExternalStoragePath(),
+                                              *UI_CAPTURE_DIRS)
+
+        if dev.PathExists(self._ui_capture_dir[dev]):
+          dev.RunShellCommand(['rm', '-rf', self._ui_capture_dir[dev]])
+        dev.RunShellCommand(['mkdir', self._ui_capture_dir[dev]])
+
       steps += [set_debug_app, edit_shared_prefs, push_test_data,
-                create_flag_changer]
+                create_flag_changer, setup_ui_capture_dir]
       if self._env.concurrent_adb:
         reraiser_thread.RunAsync(steps)
       else:
@@ -244,6 +262,20 @@ class LocalDeviceInstrumentationTestRun(
 
       valgrind_tools.SetChromeTimeoutScale(dev, None)
 
+      if self._test_instance.ui_screenshot_dir:
+        pull_ui_screen_captures(dev)
+
+    @trace_event.traced
+    def pull_ui_screen_captures(dev):
+      file_names = dev.ListDirectory(self._ui_capture_dir[dev])
+      target_path = self._test_instance.ui_screenshot_dir
+      if not os.path.exists(target_path):
+        os.makedirs(target_path)
+
+      for file_name in file_names:
+        dev.PullFile(posixpath.join(self._ui_capture_dir[dev], file_name),
+                     target_path)
+
     self._env.parallel_devices.pMap(individual_device_tear_down)
 
   def _CreateFlagChangerIfNeeded(self, device):
@@ -292,6 +324,8 @@ class LocalDeviceInstrumentationTestRun(
           device.adb, suffix='.png', dir=device.GetExternalStoragePath())
       extras[EXTRA_SCREENSHOT_FILE] = screenshot_device_file.name
 
+    extras[EXTRA_UI_CAPTURE_DIR] = self._ui_capture_dir[device]
+
     if isinstance(test, list):
       if not self._test_instance.driver_apk:
         raise Exception('driver_apk does not exist. '