Find annotated tests by exposing API in instrumentation_test_instance

build/android/pylib/instrumentation/instrumentation_test_instance.py

Index: build/android/pylib/instrumentation/instrumentation_test_instance.py
diff --git a/build/android/pylib/instrumentation/instrumentation_test_instance.py b/build/android/pylib/instrumentation/instrumentation_test_instance.py
index 12fe4309f298615ab0bfbec3e25358609875ac48..4d1f6bceeeccfab7563e19cf19b2e87a8939cef9 100644
--- a/build/android/pylib/instrumentation/instrumentation_test_instance.py
+++ b/build/android/pylib/instrumentation/instrumentation_test_instance.py
@@ -59,6 +59,10 @@ class MissingSizeAnnotationError(Exception):
         ', '.join('@' + a for a in _VALID_ANNOTATIONS))
 
 
+class ProguardPickleException(Exception):
+  pass
+
+
 # TODO(jbudorick): Make these private class methods of
 # InstrumentationTestInstance once the instrumentation test_runner is
 # deprecated.
@@ -198,6 +202,154 @@ def ParseCommandLineFlagParameters(annotations):
   return result if result else None
 
 
+def FilterTests(tests, test_filter=None, annotations=None,
+                excluded_annotations=None):
+  """Filter a list of tests
+
+  Args:
+    tests: a list of tests. e.g. [
+           {'annotations": {}, 'class': 'com.example.TestA', 'methods':[]},
+           {'annotations": {}, 'class': 'com.example.TestB', 'methods':[]}]
+    test_filter: googletest-style filter string.
+    annotations: a dict of wanted annotations for test methods.
+    exclude_annotations: a dict of annotations to exclude.
+
+  Return:
+    A list of filtered tests
+  """
+  def gtest_filter(c, m):
+    if not test_filter:
+      return True
+    # Allow fully-qualified name as well as an omitted package.
+    names = ['%s.%s' % (c['class'], m['method']),
+             '%s.%s' % (c['class'].split('.')[-1], m['method'])]
+    return unittest_util.FilterTestNames(names, test_filter)
+
+  def annotation_filter(all_annotations):
+    if not annotations:
+      return True
+    return any_annotation_matches(annotations, all_annotations)
+
+  def excluded_annotation_filter(all_annotations):
+    if not excluded_annotations:
+      return True
+    return not any_annotation_matches(excluded_annotations,
+                                      all_annotations)
+
+  def any_annotation_matches(filter_annotations, all_annotations):
+    return any(
+        ak in all_annotations
+        and annotation_value_matches(av, all_annotations[ak])
+        for ak, av in filter_annotations.iteritems())
+
+  def annotation_value_matches(filter_av, av):
+    if filter_av is None:
+      return True
+    elif isinstance(av, dict):
+      return filter_av in av['value']
+    elif isinstance(av, list):
+      return filter_av in av
+    return filter_av == av
+
+  filtered_classes = []
+  for c in tests:
+    filtered_methods = []
+    for m in c['methods']:
+      # Gtest filtering
+      if not gtest_filter(c, m):
+        continue
+
+      all_annotations = dict(c['annotations'])
+      all_annotations.update(m['annotations'])
+
+      # Enforce that all tests declare their size.
+      if not any(a in _VALID_ANNOTATIONS for a in all_annotations):
+        raise MissingSizeAnnotationError('%s.%s' % (c['class'], m['method']))
+
+      if (not annotation_filter(all_annotations)
+          or not excluded_annotation_filter(all_annotations)):
+        continue
+
+      filtered_methods.append(m)
+
+    if filtered_methods:
+      filtered_class = dict(c)
+      filtered_class['methods'] = filtered_methods
+      filtered_classes.append(filtered_class)
+
+  return filtered_classes
+
+def GetAllTests(test_jar):
+  pickle_path = '%s-proguard.pickle' % test_jar
+  try:
+    tests = _GetTestsFromPickle(pickle_path, test_jar)
+  except ProguardPickleException as e:
+    logging.info('Could not get tests from pickle: %s', e)
+    logging.info('Getting tests from JAR via proguard.')
+    tests = _GetTestsFromProguard(test_jar)
+    _SaveTestsToPickle(pickle_path, test_jar, tests)
+  return tests
+
+
+def _GetTestsFromPickle(pickle_path, jar_path):
+  if not os.path.exists(pickle_path):
+    raise ProguardPickleException('%s does not exist.' % pickle_path)
+  if os.path.getmtime(pickle_path) <= os.path.getmtime(jar_path):
+    raise ProguardPickleException(
+        '%s newer than %s.' % (jar_path, pickle_path))
+
+  with open(pickle_path, 'r') as pickle_file:
+    pickle_data = pickle.loads(pickle_file.read())
+  jar_md5 = md5sum.CalculateHostMd5Sums(jar_path)[jar_path]
+
+  if pickle_data['VERSION'] != _PICKLE_FORMAT_VERSION:
+    raise ProguardPickleException('PICKLE_FORMAT_VERSION has changed.')
+  if pickle_data['JAR_MD5SUM'] != jar_md5:
+    raise ProguardPickleException('JAR file MD5 sum differs.')
+  return pickle_data['TEST_METHODS']
+
+
+def _GetTestsFromProguard(jar_path):
+  p = proguard.Dump(jar_path)
+  class_lookup = dict((c['class'], c) for c in p['classes'])
+
+  def is_test_class(c):
+    return c['class'].endswith('Test')
+
+  def is_test_method(m):
+    return m['method'].startswith('test')
+
+  def recursive_class_annotations(c):
+    s = c['superclass']
+    if s in class_lookup:
+      a = recursive_class_annotations(class_lookup[s])
+    else:
+      a = {}
+    a.update(c['annotations'])
+    return a
+
+  def stripped_test_class(c):
+    return {
+      'class': c['class'],
+      'annotations': recursive_class_annotations(c),
+      'methods': [m for m in c['methods'] if is_test_method(m)],
+    }
+
+  return [stripped_test_class(c) for c in p['classes']
+          if is_test_class(c)]
+
+
+def _SaveTestsToPickle(pickle_path, jar_path, tests):
+  jar_md5 = md5sum.CalculateHostMd5Sums(jar_path)[jar_path]
+  pickle_data = {
+    'VERSION': _PICKLE_FORMAT_VERSION,
+    'JAR_MD5SUM': jar_md5,
+    'TEST_METHODS': tests,
+  }
+  with open(pickle_path, 'w') as pickle_file:
+    pickle.dump(pickle_data, pickle_file)
+
+
 class InstrumentationTestInstance(test_instance.TestInstance):
 
   def __init__(self, args, isolate_delegate, error_func):
@@ -294,6 +446,7 @@ class InstrumentationTestInstance(test_instance.TestInstance):
       for package_info in constants.PACKAGE_INFO.itervalues():
         if package_under_test == package_info.package:
           self._package_info = package_info
+          break
     if not self._package_info:
       logging.warning('Unable to find package info for %s', self._test_package)
 
@@ -493,144 +646,12 @@ class InstrumentationTestInstance(test_instance.TestInstance):
     return self._data_deps
 
   def GetTests(self):
-    pickle_path = '%s-proguard.pickle' % self.test_jar
-    try:
-      tests = self._GetTestsFromPickle(pickle_path, self.test_jar)
-    except self.ProguardPickleException as e:
-      logging.info('Getting tests from JAR via proguard. (%s)', str(e))
-      tests = self._GetTestsFromProguard(self.test_jar)
-      self._SaveTestsToPickle(pickle_path, self.test_jar, tests)
-    return self._ParametrizeTestsWithFlags(
-        self._InflateTests(self._FilterTests(tests)))
-
-  class ProguardPickleException(Exception):
-    pass
-
-  def _GetTestsFromPickle(self, pickle_path, jar_path):
-    if not os.path.exists(pickle_path):
-      raise self.ProguardPickleException('%s does not exist.' % pickle_path)
-    if os.path.getmtime(pickle_path) <= os.path.getmtime(jar_path):
-      raise self.ProguardPickleException(
-          '%s newer than %s.' % (jar_path, pickle_path))
-
-    with open(pickle_path, 'r') as pickle_file:
-      pickle_data = pickle.loads(pickle_file.read())
-    jar_md5 = md5sum.CalculateHostMd5Sums(jar_path)[jar_path]
-
-    try:
-      if pickle_data['VERSION'] != _PICKLE_FORMAT_VERSION:
-        raise self.ProguardPickleException('PICKLE_FORMAT_VERSION has changed.')
-      if pickle_data['JAR_MD5SUM'] != jar_md5:
-        raise self.ProguardPickleException('JAR file MD5 sum differs.')
-      return pickle_data['TEST_METHODS']
-    except TypeError as e:
-      logging.error(pickle_data)
-      raise self.ProguardPickleException(str(e))
+    tests = GetAllTests(self.test_jar)
+    filtered_tests = FilterTests(
+        tests, self._test_filter, self._annotations, self._excluded_annotations)
+    return self._ParametrizeTestsWithFlags(self._InflateTests(filtered_tests))
 
   # pylint: disable=no-self-use
-  def _GetTestsFromProguard(self, jar_path):
-    p = proguard.Dump(jar_path)
-
-    def is_test_class(c):
-      return c['class'].endswith('Test')
-
-    def is_test_method(m):
-      return m['method'].startswith('test')
-
-    class_lookup = dict((c['class'], c) for c in p['classes'])
-    def recursive_get_class_annotations(c):
-      s = c['superclass']
-      if s in class_lookup:
-        a = recursive_get_class_annotations(class_lookup[s])
-      else:
-        a = {}
-      a.update(c['annotations'])
-      return a
-
-    def stripped_test_class(c):
-      return {
-        'class': c['class'],
-        'annotations': recursive_get_class_annotations(c),
-        'methods': [m for m in c['methods'] if is_test_method(m)],
-      }
-
-    return [stripped_test_class(c) for c in p['classes']
-            if is_test_class(c)]
-
-  def _SaveTestsToPickle(self, pickle_path, jar_path, tests):
-    jar_md5 = md5sum.CalculateHostMd5Sums(jar_path)[jar_path]
-    pickle_data = {
-      'VERSION': _PICKLE_FORMAT_VERSION,
-      'JAR_MD5SUM': jar_md5,
-      'TEST_METHODS': tests,
-    }
-    with open(pickle_path, 'w') as pickle_file:
-      pickle.dump(pickle_data, pickle_file)
-
-  def _FilterTests(self, tests):
-
-    def gtest_filter(c, m):
-      if not self._test_filter:
-        return True
-      # Allow fully-qualified name as well as an omitted package.
-      names = ['%s.%s' % (c['class'], m['method']),
-               '%s.%s' % (c['class'].split('.')[-1], m['method'])]
-      return unittest_util.FilterTestNames(names, self._test_filter)
-
-    def annotation_filter(all_annotations):
-      if not self._annotations:
-        return True
-      return any_annotation_matches(self._annotations, all_annotations)
-
-    def excluded_annotation_filter(all_annotations):
-      if not self._excluded_annotations:
-        return True
-      return not any_annotation_matches(self._excluded_annotations,
-                                        all_annotations)
-
-    def any_annotation_matches(filter_annotations, all_annotations):
-      return any(
-          ak in all_annotations
-          and annotation_value_matches(av, all_annotations[ak])
-          for ak, av in filter_annotations.iteritems())
-
-    def annotation_value_matches(filter_av, av):
-      if filter_av is None:
-        return True
-      elif isinstance(av, dict):
-        return filter_av in av['value']
-      elif isinstance(av, list):
-        return filter_av in av
-      return filter_av == av
-
-    filtered_classes = []
-    for c in tests:
-      filtered_methods = []
-      for m in c['methods']:
-        # Gtest filtering
-        if not gtest_filter(c, m):
-          continue
-
-        all_annotations = dict(c['annotations'])
-        all_annotations.update(m['annotations'])
-
-        # Enforce that all tests declare their size.
-        if not any(a in _VALID_ANNOTATIONS for a in all_annotations):
-          raise MissingSizeAnnotationError('%s.%s' % (c['class'], m['method']))
-
-        if (not annotation_filter(all_annotations)
-            or not excluded_annotation_filter(all_annotations)):
-          continue
-
-        filtered_methods.append(m)
-
-      if filtered_methods:
-        filtered_class = dict(c)
-        filtered_class['methods'] = filtered_methods
-        filtered_classes.append(filtered_class)
-
-    return filtered_classes
-
   def _InflateTests(self, tests):
     inflated_tests = []
     for c in tests:
@@ -687,4 +708,3 @@ class InstrumentationTestInstance(test_instance.TestInstance):
   def TearDown(self):
     if self._isolate_delegate:
       self._isolate_delegate.Clear()
-