tools: WPR support for testing speculative_prefetch_predictor.

tools/resource_prefetch_predictor/generate_test_data.py

Index: tools/resource_prefetch_predictor/generate_test_data.py
diff --git a/tools/resource_prefetch_predictor/generate_database.py b/tools/resource_prefetch_predictor/generate_test_data.py
similarity index 72%
rename from tools/resource_prefetch_predictor/generate_database.py
rename to tools/resource_prefetch_predictor/generate_test_data.py
index 60dfd5fa629b8a7b3903307366d04e2f2db4d79b..39f8c35da5a60aae10f13881af145bfb9c9a22fd 100755
--- a/tools/resource_prefetch_predictor/generate_database.py
+++ b/tools/resource_prefetch_predictor/generate_test_data.py
@@ -6,6 +6,7 @@
 
 """Loads a set of web pages several times on a device, and extracts the
 predictor database.
+Also generates a WPR archive for another page.
 """
 
 import argparse
@@ -24,13 +25,14 @@ 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 device_setup
 from options import OPTIONS
 import page_track
 
 import prefetch_predictor_common
 
 
-_PAGE_LOAD_TIMEOUT = 20
+_PAGE_LOAD_TIMEOUT = 40
 
 
 def _CreateArgumentParser():
@@ -43,14 +45,16 @@ def _CreateArgumentParser():
                       '(one per line). URLs can be repeated.')
   parser.add_argument('--output_filename',
                       help='File to store the database in.')
+  parser.add_argument('--test_url', help='URL to record an archive of.')
+  parser.add_argument('--wpr_archive', help='WPR archive path.')
   parser.add_argument('--url_repeat',
                       help=('Number of times each URL in the input '
-                            'file is loaded.'),
-                      default=3)
+                            'file is loaded.'), default=3)
   return parser
 
 
-def _Go(chrome_controller, urls_filename, output_filename, repeats):
+def _GenerateDatabase(chrome_controller, urls_filename, output_filename,
+                      repeats):
   urls = []
   with open(urls_filename) as f:
     urls = [line.strip() for line in f.readlines()]
@@ -71,6 +75,16 @@ def _Go(chrome_controller, urls_filename, output_filename, repeats):
                   output_filename)
 
 
+def _GenerateWprArchive(device, url, archive_path):
+  with device_setup.RemoteWprHost(device, archive_path, record=True) as wpr:
+    chrome_controller = prefetch_predictor_common.Setup(
+        device, wpr.chrome_args)
+    with chrome_controller.Open() as connection:
+      page_track.PageTrack(connection)  # Registers the listeners.
+      connection.MonitorUrl(url, timeout_seconds=_PAGE_LOAD_TIMEOUT,
+                            stop_delay_multiplier=1.5)
+
+
 def main():
   devil_chromium.Initialize()
   logging.basicConfig(level=logging.INFO)
@@ -86,8 +100,9 @@ def main():
 
   chrome_controller = prefetch_predictor_common.Setup(
       device, ['--speculative-resource-prefetching=learning'])
-  _Go(chrome_controller, args.urls_filename, args.output_filename,
-      int(args.url_repeat))
+  _GenerateDatabase(chrome_controller, args.urls_filename,
+                    args.output_filename, int(args.url_repeat))
+  _GenerateWprArchive(device, args.test_url, args.wpr_archive)
 
 
 if __name__ == '__main__':