#!/usr/bin/python
#
# Copyright 2016 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.

"""Loads a set of web pages several times on a device, and extracts the
predictor database.
"""

import argparse
import logging
import os
import sys


_SRC_PATH = os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.pardir, os.pardir))
sys.path.append(os.path.join(_SRC_PATH, 'third_party', 'catapult', 'devil'))

from devil.android import device_utils

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


_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()])
  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.')
  parser.add_argument('--output_filename',
                      help='File to store the database in.')
  parser.add_argument('--url_repeat',
                      help=('Number of times each URL in the input '
                            'file is loaded.'),
                      default=3)
  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:
    urls = [line.strip() for line in f.readlines()]

  with chrome_controller.Open() as connection:
    for repeat in range(repeats):
      logging.info('Repeat #%d', repeat)
      for url in urls:
        logging.info('\tLoading %s', url)
        page_track.PageTrack(connection)  # Registers the listeners.
        connection.MonitorUrl(url, timeout_seconds=_PAGE_LOAD_TIMEOUT,
                              stop_delay_multiplier=1.5)

  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)


def main():
  logging.basicConfig(level=logging.INFO)
  parser = _CreateArgumentParser()
  args = parser.parse_args()
  OPTIONS.SetParsedArgs(args)
  devil_chromium.Initialize()
  device = _FindDevice(args.device)
  if device is None:
    logging.error('Could not find device: %s.', args.device)
    sys.exit(1)
  chrome_controller = _Setup(device)
  _Go(chrome_controller, args.urls_filename, args.output_filename,
      int(args.url_repeat))


if __name__ == '__main__':
  main()

[((19, 16, 19, 75), 'os.path.join', 'os.path.join', ({(19, 29, 19, 38): '_SRC_PATH', (19, 40, 19, 53): '"""third_party"""', (19, 55, 19, 65): '"""catapult"""', (19, 67, 19, 74): '"""devil"""'}, {}), "(_SRC_PATH, 'third_party', 'catapult', 'devil')", False, 'import os\n'), ((23, 16, 23, 59), 'os.path.join', 'os.path.join', ({(23, 29, 23, 38): '_SRC_PATH', (23, 40, 23, 47): '"""build"""', (23, 49, 23, 58): '"""android"""'}, {}), "(_SRC_PATH, 'build', 'android')", False, 'import os\n'), ((26, 16, 26, 70), 'os.path.join', 'os.path.join', ({(26, 29, 26, 38): '_SRC_PATH', (26, 40, 26, 47): '"""tools"""', (26, 49, 26, 58): '"""android"""', (26, 60, 26, 69): '"""loading"""'}, {}), "(_SRC_PATH, 'tools', 'android', 'loading')", False, 'import os\n'), ((55, 12, 55, 53), 'devil.android.device_utils.DeviceUtils.HealthyDevices', 'device_utils.DeviceUtils.HealthyDevices', ({}, {}), '()', False, 'from devil.android import device_utils\n'), ((66, 22, 66, 63), 'controller.RemoteChromeController', 'controller.RemoteChromeController', ({(66, 56, 66, 62): 'device'}, {}), '(device)', False, 'import controller\n'), ((97, 2, 97, 41), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((100, 2, 100, 29), 'options.OPTIONS.SetParsedArgs', 'OPTIONS.SetParsedArgs', ({(100, 24, 100, 28): 'args'}, {}), '(args)', False, 'from options import OPTIONS\n'), ((101, 2, 101, 29), 'devil_chromium.Initialize', 'devil_chromium.Initialize', ({}, {}), '()', False, 'import devil_chromium\n'), ((18, 4, 18, 29), 'os.path.dirname', 'os.path.dirname', ({(18, 20, 18, 28): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((104, 4, 104, 60), 'logging.error', 'logging.error', ({(104, 18, 104, 46): '"""Could not find device: %s."""', (104, 48, 104, 59): 'args.device'}, {}), "('Could not find device: %s.', args.device)", False, 'import logging\n'), ((105, 4, 105, 15), 'sys.exit', 'sys.exit', ({(105, 13, 105, 14): '(1)'}, {}), '(1)', False, 'import sys\n'), ((67, 19, 67, 42), 'options.OPTIONS.ChromePackage', 'OPTIONS.ChromePackage', ({}, {}), '()', False, 'from options import OPTIONS\n'), ((81, 6, 81, 40), 'logging.info', 'logging.info', ({(81, 19, 81, 31): '"""Repeat #%d"""', (81, 33, 81, 39): 'repeat'}, {}), "('Repeat #%d', repeat)", False, 'import logging\n'), ((89, 19, 89, 42), 'options.OPTIONS.ChromePackage', 'OPTIONS.ChromePackage', ({}, {}), '()', False, 'from options import OPTIONS\n'), ((92, 6, 92, 29), 'options.OPTIONS.ChromePackage', 'OPTIONS.ChromePackage', ({}, {}), '()', False, 'from options import OPTIONS\n'), ((40, 15, 40, 40), 'options.OPTIONS.GetParentParser', 'OPTIONS.GetParentParser', ({}, {}), '()', False, 'from options import OPTIONS\n'), ((83, 8, 83, 41), 'logging.info', 'logging.info', ({(83, 21, 83, 35): '"""\tLoading %s"""', (83, 37, 83, 40): 'url'}, {}), "('\\tLoading %s', url)", False, 'import logging\n'), ((84, 8, 84, 40), 'page_track.PageTrack', 'page_track.PageTrack', ({(84, 29, 84, 39): 'connection'}, {}), '(connection)', False, 'import page_track\n')]