Support for capturing screen after each startup run

Change-Id: Id04fcdbba76897bc71570c2b4c2ab8f2ddd2c399
diff --git a/tools/startup/adb_utils.py b/tools/startup/adb_utils.py
index 2369aa7..a5a4e89 100755
--- a/tools/startup/adb_utils.py
+++ b/tools/startup/adb_utils.py
@@ -56,6 +56,13 @@
   cmd = create_adb_cmd('shell am broadcast -a %s %s' % (action, component), device_id)
   return subprocess.check_output(cmd).decode('utf-8').strip().splitlines()
 
+def capture_screen(target, device_id=None):
+  print('Taking screenshot to %s' % target)
+  tmp = '/sdcard/screencap.png'
+  cmd = create_adb_cmd('shell screencap -p %s' % tmp, device_id)
+  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
+  pull(tmp, target, device_id)
+
 def create_adb_cmd(arguments, device_id=None):
   assert isinstance(arguments, list) or isinstance(arguments, str)
   cmd = ['adb']
@@ -135,8 +142,7 @@
   with utils.TempDir() as temp:
     source = get_profile_path(app_id)
     target = os.path.join(temp, 'primary.prof')
-    cmd = create_adb_cmd('pull %s %s' % (source, target), device_id)
-    subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
+    pull(source, target, device_id)
     with open(target, 'rb') as f:
       return f.read()
 
@@ -288,6 +294,10 @@
   }
   return teardown_options
 
+def pull(source, target, device_id=None):
+  cmd = create_adb_cmd('pull %s %s' % (source, target), device_id)
+  subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
+
 def root(device_id=None):
   cmd = create_adb_cmd('root', device_id)
   subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
@@ -364,6 +374,8 @@
 
 def parse_options(argv):
   result = argparse.ArgumentParser(description='Run adb utils.')
+  result.add_argument('--capture-screen',
+                      help='Capture screen to given file')
   result.add_argument('--device-id',
                       help='Device id (e.g., emulator-5554).')
   result.add_argument('--device-pin',
@@ -385,6 +397,8 @@
 
 def main(argv):
   (options, args) = parse_options(argv)
+  if options.capture_screen:
+    capture_screen(options.capture_screen, options.device_id)
   if options.ensure_screen_off:
     ensure_screen_off(options.device_id)
   elif options.get_screen_state:
diff --git a/tools/startup/measure_startup.py b/tools/startup/measure_startup.py
index 04076ec..46609fe 100755
--- a/tools/startup/measure_startup.py
+++ b/tools/startup/measure_startup.py
@@ -55,7 +55,7 @@
     out_dir = os.path.join(options.out_dir, str(iteration))
     teardown_options = setup_for_run(apk, out_dir, options)
     data = run(out_dir, options, tmp_dir)
-    teardown_for_run(options, teardown_options)
+    teardown_for_run(out_dir, options, teardown_options)
     add_data(data_total, data)
     print('Result:')
     print(data)
@@ -119,9 +119,13 @@
   adb_utils.drop_caches(options.device_id)
   return teardown_options
 
-def teardown_for_run(options, teardown_options):
+def teardown_for_run(out_dir, options, teardown_options):
   assert adb_utils.get_screen_state(options.device_id).is_on_and_unlocked()
 
+  if options.capture_screen:
+    target = os.path.join(out_dir, 'screen.png')
+    adb_utils.capture_screen(target, options.device_id)
+
   if options.cooldown > 0:
     adb_utils.teardown_after_interaction_with_device(
         teardown_options, options.device_id)
@@ -248,6 +252,10 @@
   result.add_argument('--apk',
                       help='Path to the APK',
                       required=True)
+  result.add_argument('--capture-screen',
+                      help='Take a screenshot after each test',
+                      default=False,
+                      action='store_true')
   result.add_argument('--cooldown',
                       help='Seconds to wait before running each iteration',
                       default=0,