tools: Local tests for the speculative prefetch predictor.

tools/resource_prefetch_predictor/generate_database.py

Index: tools/resource_prefetch_predictor/generate_database.py
diff --git a/tools/resource_prefetch_predictor/generate_database.py b/tools/resource_prefetch_predictor/generate_database.py
index 9c3aa677e610f9e2bf81b41d5bae0ca83fbbae6f..60dfd5fa629b8a7b3903307366d04e2f2db4d79b 100755
--- a/tools/resource_prefetch_predictor/generate_database.py
+++ b/tools/resource_prefetch_predictor/generate_database.py
@@ -12,7 +12,7 @@ import argparse
 import logging
 import os
 import sys
-
+import time
 
 _SRC_PATH = os.path.abspath(os.path.join(
     os.path.dirname(__file__), os.pardir, os.pardir))
@@ -24,10 +24,11 @@ sys.path.append(os.path.join(_SRC_PATH, 'build', 'android'))
 import devil_chromium
 
 sys.path.append(os.path.join(_SRC_PATH, 'tools', 'android', 'loading'))
-import controller
 from options import OPTIONS
 import page_track
 
+import prefetch_predictor_common
+
 
 _PAGE_LOAD_TIMEOUT = 20
 
@@ -35,9 +36,8 @@ _PAGE_LOAD_TIMEOUT = 20
 def _CreateArgumentParser():
   """Creates and returns the argument parser."""
   parser = argparse.ArgumentParser(
-      description=('Loads a set of web pages several times on a device, and '
-                   'extracts the predictor database.'),
-      parents=[OPTIONS.GetParentParser()])
+      ('Loads a set of web pages several times on a device, and extracts the '
+       'predictor database.'), parents=[OPTIONS.GetParentParser()])
   parser.add_argument('--device', help='Device ID')
   parser.add_argument('--urls_filename', help='File containing a list of URLs '
                       '(one per line). URLs can be repeated.')
@@ -50,27 +50,6 @@ def _CreateArgumentParser():
   return parser
 
 
-def _FindDevice(device_id):
-  """Returns a device matching |device_id| or the first one if None, or None."""
-  devices = device_utils.DeviceUtils.HealthyDevices()
-  if device_id is None:
-    return devices[0]
-  matching_devices = [d for d in devices if str(d) == device_id]
-  if not matching_devices:
-    return None
-  return matching_devices[0]
-
-
-def _Setup(device):
-  """Sets up a device and returns an instance of RemoteChromeController."""
-  chrome_controller = controller.RemoteChromeController(device)
-  device.ForceStop(OPTIONS.ChromePackage().package)
-  chrome_controller.AddChromeArguments(
-      ['--speculative-resource-prefetching=learning'])
-  chrome_controller.ResetBrowserState()
-  return chrome_controller
-
-
 def _Go(chrome_controller, urls_filename, output_filename, repeats):
   urls = []
   with open(urls_filename) as f:
@@ -84,26 +63,29 @@ def _Go(chrome_controller, urls_filename, output_filename, repeats):
         page_track.PageTrack(connection)  # Registers the listeners.
         connection.MonitorUrl(url, timeout_seconds=_PAGE_LOAD_TIMEOUT,
                               stop_delay_multiplier=1.5)
+        time.sleep(2)  # Reduces flakiness.
 
   device = chrome_controller.GetDevice()
   device.ForceStop(OPTIONS.ChromePackage().package)
-  database_filename = (
-      '/data/user/0/%s/app_chrome/Default/Network Action Predictor' %
-      OPTIONS.ChromePackage().package)
-  device.PullFile(database_filename, output_filename)
+  device.PullFile(prefetch_predictor_common.DatabaseDevicePath(),
+                  output_filename)
 
 
 def main():
+  devil_chromium.Initialize()
   logging.basicConfig(level=logging.INFO)
+
   parser = _CreateArgumentParser()
   args = parser.parse_args()
   OPTIONS.SetParsedArgs(args)
-  devil_chromium.Initialize()
-  device = _FindDevice(args.device)
+
+  device = prefetch_predictor_common.FindDevice(args.device)
   if device is None:
     logging.error('Could not find device: %s.', args.device)
     sys.exit(1)
-  chrome_controller = _Setup(device)
+
+  chrome_controller = prefetch_predictor_common.Setup(
+      device, ['--speculative-resource-prefetching=learning'])
   _Go(chrome_controller, args.urls_filename, args.output_filename,
       int(args.url_repeat))