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 cd8617cfc15f7bb84fa79a60caa1a71261dd2de9..4529c175fefa19f255adf63abed5eccf842d21e1 100644
--- a/build/android/pylib/instrumentation/instrumentation_test_instance.py
+++ b/build/android/pylib/instrumentation/instrumentation_test_instance.py
@@ -187,6 +187,145 @@ def ParseCommandLineFlagParameters(annotations):
         result.append(ParamsTuple(to_add, to_remove))
   return result if result else None
 
+def GetTestsFromPickle(pickle_path, jar_path):

[jbudorick, 2016/04/06 20:33:30]

Thanks for putting up with my waffling on this.

Let's go with this version, but make everything that isn't called by your
function module-private (i.e., have the names start with _)

[Yoland Yan(Google), 2016/04/06 22:04:19]

dude, ofc
Done.

+  """Doc stub"""
+  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]
+
+  try:
+    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']
+  except TypeError as e:
+    logging.error(pickle_data)
+    raise ProguardPickleException(str(e))
+
+# pylint: disable=no-self-use
+def GetTestsFromProguard(jar_path):
+  """Doc stub"""
+  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(pickle_path, jar_path, tests):
+  """Doc stub"""
+  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 TotalTestAmount(test_jar):
+  """Doc stub"""
+  total_test_amount = 0
+  tests = _GetAllTests(test_jar)
+  for c in tests:
+    total_test_amount += len(c['methods'])
+  return total_test_amount
+
+def GetFilteredTests(test_jar, test_filter, annotations, exclude_annotations):
+  """Doc stub"""
+  tests = _GetAllTests(test_jar)
+  filtered_tests = FilterTests(tests, test_filter, annotations,
+                               exclude_annotations)
+  return filtered_tests
+
+def FilterTests(tests, test_filter, annotations, excluded_annotations):
+  """Doc stub"""
+
+  def gtest_filter(c, m):
+    t = ['%s.%s' % (c['class'].split('.')[-1], m['method'])]
+    return (not test_filter
+            or unittest_util.FilterTestNames(t, 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(annotations, all_annotations):
+    return any(
+        ak in all_annotations and (av is None or av == all_annotations[ak])
+        for ak, av in annotations.iteritems())
+
+  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'])
+      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):
+  """Doc stub"""
+  pickle_path = '%s-proguard.pickle' % test_jar
+  try:
+    tests = GetTestsFromPickle(pickle_path, test_jar)
+  except ProguardPickleException as e:
+    logging.info('Getting tests from JAR via proguard. (%s)', str(e))
+    tests = GetTestsFromProguard(test_jar)
+    SaveTestsToPickle(pickle_path, test_jar, tests)
+  return tests
+
+class ProguardPickleException(Exception):
+  pass
 
 class InstrumentationTestInstance(test_instance.TestInstance):
 
@@ -465,125 +604,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))
-
-  # 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):
-      t = ['%s.%s' % (c['class'].split('.')[-1], m['method'])]
-      return (not self._test_filter
-              or unittest_util.FilterTestNames(t, 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(annotations, all_annotations):
-      return any(
-          ak in all_annotations and (av is None or av == all_annotations[ak])
-          for ak, av in annotations.iteritems())
-
-    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'])
-        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
+    filtered_tests = GetFilteredTests(
+        self.test_jar,
+        self._test_filter,
+        self._annotations,
+        self._excluded_annotations)
+    return self._ParametrizeTestsWithFlags(self._InflateTests(filtered_tests))
 
   def _InflateTests(self, tests):
     inflated_tests = []
@@ -641,4 +667,3 @@ class InstrumentationTestInstance(test_instance.TestInstance):
   def TearDown(self):
     if self._isolate_delegate:
       self._isolate_delegate.Clear()
-