Move screenshot capture to Java-side.

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

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import logging
 import os
 import posixpath
 import re
 import time
 
 from devil.android import device_errors
+from devil.android import device_temp_file
 from devil.android import flag_changer
 from devil.android.sdk import shared_prefs
 from devil.utils import reraiser_thread
 from pylib import valgrind_tools
 from pylib.android import logdog_logcat_monitor
 from pylib.base import base_test_result
 from pylib.instrumentation import instrumentation_test_instance
 from pylib.local.device import local_device_environment
 from pylib.local.device import local_device_test_run
 from pylib.utils import google_storage_helper
@@ skipping 103 matching lines @@
       def set_debug_app():
         # Set debug app in order to enable reading command line flags on user
         # builds
         if self._test_instance.flags:
           if not self._test_instance.package_info:
             logging.error("Couldn't set debug app: no package info")
           elif not self._test_instance.package_info.package:
             logging.error("Couldn't set debug app: no package defined")
           else:
             dev.RunShellCommand(['am', 'set-debug-app', '--persistent',
-                                  self._test_instance.package_info.package],
+                                 self._test_instance.package_info.package],
                                 check_return=True)
       @trace_event.traced
       def edit_shared_prefs():
         for pref in self._test_instance.edit_shared_prefs:
           prefs = shared_prefs.SharedPrefs(dev, pref['package'],
                                            pref['filename'])
           prefs.Load()
           for key in pref.get('remove', []):
             try:
               prefs.Remove(key)
@@ skipping 93 matching lines @@
     return tests
 
   #override
   def _GetUniqueTestName(self, test):
     return instrumentation_test_instance.GetUniqueTestName(test)
 
   #override
   def _RunTest(self, device, test):
     extras = {}
 
-    flags = None
+    flags_to_add = []
+    flags_to_remove = []
     test_timeout_scale = None
     if self._test_instance.coverage_directory:
       coverage_basename = '%s.ec' % ('%s_group' % test[0]['method']
           if isinstance(test, list) else test['method'])
       extras['coverage'] = 'true'
       coverage_directory = os.path.join(
           device.GetExternalStoragePath(), 'chrome', 'test', 'coverage')
       coverage_device_file = os.path.join(
           coverage_directory, coverage_basename)
       extras['coverageFile'] = coverage_device_file
+    # Save screenshot if screenshot dir is specified (save locally) or if
+    # a GS bucket is passed (save in cloud).
+    screenshot_device_file = None

[jbudorick, 2017/05/08 22:32:06]

optional nit: could do something here similar to what we do in
local_device_gtest_run w/ a null context manager. I need to move that up to
contextlib_ext, though, so I'm ok with you holding off & landing this.

+    if (self._test_instance.screenshot_dir or
+        self._test_instance.gs_results_bucket):
+      screenshot_device_file = device_temp_file.DeviceTempFile(
+          device.adb, suffix='.png', dir=device.GetExternalStoragePath())
+      flags_to_add.append('--screenshot-file=%s' % screenshot_device_file.name)
 
     if isinstance(test, list):
       if not self._test_instance.driver_apk:
         raise Exception('driver_apk does not exist. '
                         'Please build it and try again.')
       if any(t.get('is_junit4') for t in test):
         raise Exception('driver apk does not support JUnit4 tests')
 
       def name_and_timeout(t):
         n = instrumentation_test_instance.GetTestName(t)
@@ skipping 16 matching lines @@
       test_display_name = self._GetUniqueTestName(test)
       if test['is_junit4']:
         target = '%s/%s' % (
             self._test_instance.test_package,
             self._test_instance.test_runner_junit4)
       else:
         target = '%s/%s' % (
             self._test_instance.test_package, self._test_instance.test_runner)
       extras['class'] = test_name
       if 'flags' in test:
-        flags = test['flags']
+        flags_to_add.extend(test['flags'].add)
+        flags_to_remove.extend(test['flags'].remove)
       timeout = self._GetTimeoutFromAnnotations(
         test['annotations'], test_display_name)
 
       test_timeout_scale = self._GetTimeoutScaleFromAnnotations(
           test['annotations'])
       if test_timeout_scale and test_timeout_scale != 1:
         valgrind_tools.SetChromeTimeoutScale(
             device, test_timeout_scale * self._test_instance.timeout_scale)
 
     logging.info('preparing to run %s: %s', test_display_name, test)
 
-    if flags:
+    if flags_to_add or flags_to_remove:
       self._CreateFlagChangerIfNeeded(device)
       self._flag_changers[str(device)].PushFlags(
-        add=flags.add, remove=flags.remove)
+        add=flags_to_add, remove=flags_to_remove)
 
     try:
       device.RunShellCommand(
           ['log', '-p', 'i', '-t', _TAG, 'START %s' % test_name],
           check_return=True)
       time_ms = lambda: int(time.time() * 1e3)
       start_ms = time_ms()
 
       stream_name = 'logcat_%s_%s_%s' % (
           test_name.replace('#', '.'),
           time.strftime('%Y%m%dT%H%M%S-UTC', time.gmtime()),
           device.serial)
       logmon = logdog_logcat_monitor.LogdogLogcatMonitor(
           device.adb, stream_name, filter_specs=LOGCAT_FILTERS)
 
       with contextlib_ext.Optional(
           logmon, self._test_instance.should_save_logcat):
         with contextlib_ext.Optional(
             trace_event.trace(test_name),
             self._env.trace_output):
           output = device.StartInstrumentation(
               target, raw=True, extras=extras, timeout=timeout, retries=0)
       logcat_url = logmon.GetLogcatURL()
     finally:
       device.RunShellCommand(
           ['log', '-p', 'i', '-t', _TAG, 'END %s' % test_name],
           check_return=True)
       duration_ms = time_ms() - start_ms
-      if flags:
+      if flags_to_add or flags_to_remove:
         self._flag_changers[str(device)].Restore()
       if test_timeout_scale:
         valgrind_tools.SetChromeTimeoutScale(
             device, self._test_instance.timeout_scale)
 
     # TODO(jbudorick): Make instrumentation tests output a JSON so this
     # doesn't have to parse the output.
     result_code, result_bundle, statuses = (
         self._test_instance.ParseAmInstrumentRawOutput(output))
     results = self._test_instance.GenerateTestResults(
         result_code, result_bundle, statuses, start_ms, duration_ms)
     for result in results:
       if logcat_url:
         result.SetLink('logcat', logcat_url)
 
     # Update the result name if the test used flags.
-    if flags:
+    if flags_to_add or flags_to_remove:
       for r in results:
         if r.GetName() == test_name:
           r.SetName(test_display_name)
 
     # Add UNKNOWN results for any missing tests.
     iterable_test = test if isinstance(test, list) else [test]
     test_names = set(self._GetUniqueTestName(t) for t in iterable_test)
     results_names = set(r.GetName() for r in results)
     results.extend(
         base_test_result.BaseTestResult(u, base_test_result.ResultType.UNKNOWN)
@@ skipping 11 matching lines @@
     #   - clearing the application state while persisting permissions
     if any(r.GetType() not in (base_test_result.ResultType.PASS,
                                base_test_result.ResultType.SKIP)
            for r in results):
       with contextlib_ext.Optional(
           tempfile_ext.NamedTemporaryDirectory(),
           self._test_instance.screenshot_dir is None and
               self._test_instance.gs_results_bucket) as screenshot_host_dir:
         screenshot_host_dir = (
             self._test_instance.screenshot_dir or screenshot_host_dir)
-        if screenshot_host_dir:
-          file_name = '%s-%s.png' % (
-              test_display_name,
-              time.strftime('%Y%m%dT%H%M%S-UTC', time.gmtime()))
-          screenshot_file = device.TakeScreenshot(
-              os.path.join(screenshot_host_dir, file_name))
-          logging.info(
-              'Saved screenshot for %s to %s.',
-              test_display_name, screenshot_file)
-          if self._test_instance.gs_results_bucket:
-            link = google_storage_helper.upload(
-                google_storage_helper.unique_name('screenshot', device=device),
-                screenshot_file,
-                bucket=self._test_instance.gs_results_bucket + '/screenshots')
-            for result in results:
-              result.SetLink('post_test_screenshot', link)
+        self._SaveScreenshot(device, screenshot_host_dir,
+                             screenshot_device_file, test_display_name,
+                             results)
 
       logging.info('detected failure in %s. raw output:', test_display_name)
       for l in output:
         logging.info('  %s', l)
       if (not self._env.skip_clear_data
           and self._test_instance.package_info):
         permissions = (
             self._test_instance.apk_under_test.GetPermissions()
             if self._test_instance.apk_under_test
             else None)
@@ skipping 20 matching lines @@
                 include_stack_symbols=False,
                 wipe_tombstones=True)
             stream_name = 'tombstones_%s_%s' % (
                 time.strftime('%Y%m%dT%H%M%S-UTC', time.gmtime()),
                 device.serial)
             tombstones_url = logdog_helper.text(
                 stream_name, '\n'.join(resolved_tombstones))
           result.SetLink('tombstones', tombstones_url)
     return results, None
 
+  def _SaveScreenshot(self, device, screenshot_host_dir, screenshot_device_file,
+                      test_name, results):
+    if screenshot_host_dir:
+      screenshot_host_file = os.path.join(
+          screenshot_host_dir,
+          '%s-%s.png' % (
+              test_name,
+              time.strftime('%Y%m%dT%H%M%S-UTC', time.gmtime())))
+      if device.FileExists(screenshot_device_file.name):
+        try:
+          device.PullFile(screenshot_device_file.name, screenshot_host_file)
+        finally:
+          screenshot_device_file.close()
+
+        logging.info(
+            'Saved screenshot for %s to %s.',
+            test_name, screenshot_host_file)
+        if self._test_instance.gs_results_bucket:
+          link = google_storage_helper.upload(
+              google_storage_helper.unique_name(
+                  'screenshot', device=device),
+              screenshot_host_file,
+              bucket=('%s/screenshots' %
+                      self._test_instance.gs_results_bucket))
+          for result in results:
+            result.SetLink('post_test_screenshot', link)
+
   #override
   def _ShouldRetry(self, test):
     if 'RetryOnFailure' in test.get('annotations', {}):
       return True
 
     # TODO(jbudorick): Remove this log message once @RetryOnFailure has been
     # enabled for a while. See crbug.com/619055 for more details.
     logging.error('Default retries are being phased out. crbug.com/619055')
     return False
 
@@ skipping 15 matching lines @@
       if k in annotations:
         timeout = v
         break
     else:
       logging.warning('Using default 1 minute timeout for %s', test_name)
       timeout = 60
 
     timeout *= cls._GetTimeoutScaleFromAnnotations(annotations)
 
     return timeout