[telemetry] Prototype for specifying page sets using Python imports and Telemetry discovery.

tools/telemetry/telemetry/test_runner.py

Index: tools/telemetry/telemetry/test_runner.py
diff --git a/tools/telemetry/telemetry/test_runner.py b/tools/telemetry/telemetry/test_runner.py
index 66b923bb192c7c266f4bd16df4a880e92b9e110f..3799647ec1b4be51e0eba82994c974bc49d642a9 100644
--- a/tools/telemetry/telemetry/test_runner.py
+++ b/tools/telemetry/telemetry/test_runner.py
@@ -17,14 +17,41 @@ from telemetry import test
 from telemetry.core import browser_options
 from telemetry.core import command_line
 from telemetry.core import discover
-from telemetry.core import environment
-from telemetry.core import util
 from telemetry.page import page_set
 from telemetry.page import page_test
 from telemetry.page import profile_creator
 from telemetry.util import find_dependencies
 
 
+class Environment(object):
+  def __init__(self, **kwargs):
+    self._top_level_dir = kwargs['top_level_dir']
+    self._test_dirs = kwargs.get('test_dirs', [])
+    self._page_test_dirs = kwargs.get('page_test_dirs', [])
+    self._page_set_dirs = kwargs.get('page_set_dirs', [])
+    self._test_aliases = kwargs.get('test_aliases', None)
+
+  @property
+  def top_level_dir(self):
+    return self._top_level_dir
+
+  @property
+  def test_dirs(self):
+    return self._test_dirs
+
+  @property
+  def page_test_dirs(self):
+    return self._page_test_dirs
+
+  @property
+  def page_set_dirs(self):
+    return self._page_set_dirs
+
+  @property
+  def test_aliases(self):
+    return self._test_aliases
+
+
 class Deps(find_dependencies.FindDependenciesCommand):
   """Prints all dependencies"""
 
@@ -52,15 +79,17 @@ class List(command_line.OptparseCommand):
   usage = '[test_name] [<options>]'
 
   @classmethod
-  def AddCommandLineArgs(cls, parser):
+  def AddCommandLineArgs(cls, parser, _):  # pylint: disable=W0221
     parser.add_option('-j', '--json', action='store_true')
 
   @classmethod
-  def ProcessCommandLineArgs(cls, parser, args):
+  def ProcessCommandLineArgs(cls, parser, args, environment):
+    # pylint: disable=W0221
     if not args.positional_args:
-      args.tests = _Tests()
+      args.tests = _Tests(environment)
     elif len(args.positional_args) == 1:
-      args.tests = _MatchTestName(args.positional_args[0], exact_matches=False)
+      args.tests = _MatchInput(args.positional_args[0], _Tests(environment),
+                               exact_matches=False)
     else:
       parser.error('Must provide at most one test name.')
 
@@ -91,13 +120,14 @@ class Run(command_line.OptparseCommand):
     return parser
 
   @classmethod
-  def AddCommandLineArgs(cls, parser):
+  def AddCommandLineArgs(cls, parser, environment):  # pylint: disable=W0221
     test.AddCommandLineArgs(parser)
 
     # Allow tests to add their own command line options.
     matching_tests = []
     for arg in sys.argv[1:]:
-      matching_tests += _MatchTestName(arg)
+      matching_tests += _MatchInput(arg, _Tests(environment),
+                                    environment.test_aliases)
 
     if matching_tests:
       # TODO(dtu): After move to argparse, add command-line args for all tests
@@ -108,45 +138,63 @@ class Run(command_line.OptparseCommand):
       matching_test.SetArgumentDefaults(parser)
 
   @classmethod
-  def ProcessCommandLineArgs(cls, parser, args):
+  def ProcessCommandLineArgs(cls, parser, args, environment):
+    # pylint: disable=W0221
     if not args.positional_args:
-      _PrintTestList(_Tests())
+      print >> sys.stderr, 'Need to specify a test.'
+      print >> sys.stderr
+      _PrintTestList(_Tests(environment))
       sys.exit(-1)
 
-    input_test_name = args.positional_args[0]
-    matching_tests = _MatchTestName(input_test_name)
+    input_test = args.positional_args[0]
+    matching_tests = _MatchInput(input_test, _Tests(environment),
+                                 environment.test_aliases)
     if not matching_tests:
-      print >> sys.stderr, 'No test named "%s".' % input_test_name
+      print >> sys.stderr, 'No test named "%s".' % input_test
       print >> sys.stderr
-      _PrintTestList(_Tests())
+      _PrintTestList(_Tests(environment))
       sys.exit(-1)
 
     if len(matching_tests) > 1:
-      print >> sys.stderr, 'Multiple tests named "%s".' % input_test_name
+      print >> sys.stderr, 'Multiple tests named "%s".' % input_test
       print >> sys.stderr, 'Did you mean one of these?'
       print >> sys.stderr
       _PrintTestList(matching_tests)
       sys.exit(-1)
 
     test_class = matching_tests.pop()
+
     if issubclass(test_class, page_test.PageTest):
       if len(args.positional_args) < 2:
-        parser.error('Need to specify a page set for "%s".' % test_class.Name())
+        print >> sys.stderr, 'Need to specify a page set for "%s".' % (
+            test_class.Name())
+        print >> sys.stderr
+        _PrintPageSetList(_PageSets(environment))
+        sys.exit(-1)
+
       if len(args.positional_args) > 2:
         parser.error('Too many arguments.')
-      page_set_path = args.positional_args[1]
-      if not os.path.exists(page_set_path):
-        parser.error('Page set not found.')
-      if not (os.path.isfile(page_set_path) and
-              discover.IsPageSetFile(page_set_path)):
-        parser.error('Unsupported page set file format.')
+
+      input_page_set = args.positional_args[1]
+      matching_page_sets = _MatchInput(input_page_set, _PageSets(environment))
+      if not matching_page_sets:
+        print >> sys.stderr, 'No page_set named "%s".' % input_page_set
+        print >> sys.stderr
+        _PrintPageSetList(_PageSets(environment))
+        sys.exit(-1)
+
+      if len(matching_page_sets) > 1:
+        print >> sys.stderr, 'Multiple page_sets named "%s".' % input_page_set
+        print >> sys.stderr, 'Did you mean one of these?'
+        print >> sys.stderr
+        _PrintPageSetList(matching_page_sets)
+        sys.exit(-1)
+
+      page_set_class = matching_page_sets.pop()
 
       class TestWrapper(test.Test):
         test = test_class
-
-        @classmethod
-        def CreatePageSet(cls, options):
-          return page_set.PageSet.FromFile(page_set_path)
+        page_set = page_set_class
 
       test_class = TestWrapper
     else:
@@ -179,12 +227,14 @@ def _Commands():
 
 
 @decorators.Cache
-def _Tests():
+def _Tests(environment):
   tests = []
-  for base_dir in config.base_paths:
-    tests += discover.DiscoverClasses(base_dir, base_dir, test.Test,
+  for search_dir in environment.test_dirs:
+    tests += discover.DiscoverClasses(search_dir, environment.top_level_dir,
+                                      test.Test,
                                       index_by_class_name=True).values()
-    page_tests = discover.DiscoverClasses(base_dir, base_dir,
+  for search_dir in environment.page_test_dirs:
+    page_tests = discover.DiscoverClasses(search_dir, environment.top_level_dir,
                                           page_test.PageTest,
                                           index_by_class_name=True).values()
     tests += [test_class for test_class in page_tests
@@ -192,30 +242,39 @@ def _Tests():
   return tests
 
 
-def _MatchTestName(input_test_name, exact_matches=True):
-  def _Matches(input_string, search_string):
-    if search_string.startswith(input_string):
+@decorators.Cache
+def _PageSets(environment):
+  page_sets = []
+  for search_dir in environment.page_set_dirs:
+    page_sets += discover.DiscoverClasses(search_dir, environment.top_level_dir,
+                                          page_set.PageSet,
+                                          index_by_class_name=True).values()
+  return page_sets
+
+
+def _MatchInput(query, library, aliases=None, exact_matches=True):
+  def _Matches(query, candidate):
+    if candidate.startswith(query):
       return True
-    for part in search_string.split('.'):
-      if part.startswith(input_string):
+    for part in candidate.split('.'):
+      if part.startswith(query):
         return True
     return False
 
   # Exact matching.
   if exact_matches:
     # Don't add aliases to search dict, only allow exact matching for them.
-    if input_test_name in config.test_aliases:
-      exact_match = config.test_aliases[input_test_name]
+    if aliases and query in aliases:
+      exact_match = aliases[query]
     else:
-      exact_match = input_test_name
+      exact_match = query
 
-    for test_class in _Tests():
-      if exact_match == test_class.Name():
-        return [test_class]
+    for cls in library:
+      if exact_match == cls.Name():
+        return [cls]
 
   # Fuzzy matching.
-  return [test_class for test_class in _Tests()
-          if _Matches(input_test_name, test_class.Name())]
+  return [cls for cls in library if _Matches(query, cls.Name())]
 
 
 def _PrintTestList(tests):
@@ -245,10 +304,22 @@ def _PrintTestList(tests):
     print >> sys.stderr
 
 
-config = environment.Environment([util.GetBaseDir()])
+def _PrintPageSetList(page_sets):
+  if not page_sets:
+    print >> sys.stderr, 'No page sets found!'
+    return
+
+  # Align the page set names to the longest one.
+  format_string = '  %%-%ds %%s' % max(len(t.Name()) for t in page_sets)
+
+  print >> sys.stderr, 'Available page sets are:'
+  for page_set_class in sorted(page_sets, key=lambda t: t.Name()):
+    print >> sys.stderr, format_string % (
+        page_set_class.Name(), page_set_class.Description())
+  print >> sys.stderr
 
 
-def main():
+def main(environment):
   # Get the command name from the command line.
   if len(sys.argv) > 1 and sys.argv[1] == '--help':
     sys.argv[1] = 'help'
@@ -276,10 +347,10 @@ def main():
 
   # Parse and run the command.
   parser = command.CreateParser()
-  command.AddCommandLineArgs(parser)
+  command.AddCommandLineArgs(parser, environment)
   options, args = parser.parse_args()
   if commands:
     args = args[1:]
   options.positional_args = args
-  command.ProcessCommandLineArgs(parser, options)
+  command.ProcessCommandLineArgs(parser, options, environment)
   return command().Run(options)