Find annotated tests by exposing API in instrumentation_test_instance

build/android/pylib/instrumentation/instrumentation_test_instance.py

 # Copyright 2015 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.
 
 import collections
 import copy
 import logging
 import os
 import pickle
 import re
@@ skipping 41 matching lines @@
 _PICKLE_FORMAT_VERSION = 10
 
 
 class MissingSizeAnnotationError(Exception):
   def __init__(self, class_name):
     super(MissingSizeAnnotationError, self).__init__(class_name +
         ': Test method is missing required size annotation. Add one of: ' +
         ', '.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.
 def ParseAmInstrumentRawOutput(raw_output):
   """Parses the output of an |am instrument -r| call.
 
   Args:
     raw_output: the output of an |am instrument -r| call as a list of lines
   Returns:
     A 3-tuple containing:
@@ skipping 119 matching lines @@
         to_remove = []
         for a in p.get('arguments', []):
           if a['name'] == 'add':
             to_add = ['--%s' % f for f in a['stringArray']]
           elif a['name'] == 'remove':
             to_remove = ['--%s' % f for f in a['stringArray']]
         result.append(ParamsTuple(to_add, to_remove))
   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(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'])
+
+      # 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):
     super(InstrumentationTestInstance, self).__init__()
 
     self._additional_apks = []
     self._apk_under_test = None
     self._apk_under_test_incremental_install_script = None
     self._package_info = None
     self._suite = None
@@ skipping 76 matching lines @@
 
     self._test_package = self._test_apk.GetPackageName()
     self._test_runner = self._test_apk.GetInstrumentationName()
 
     self._package_info = None
     if self._apk_under_test:
       package_under_test = self._apk_under_test.GetPackageName()
       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)
 
     for apk in args.additional_apks:
       if not os.path.exists(apk):
         error_func('Unable to find additional APK: %s' % apk)
     self._additional_apks = (
         [apk_helper.ToHelper(x) for x in args.additional_apks])
 
   def _initializeDataDependencyAttributes(self, args, isolate_delegate):
@@ skipping 179 matching lines @@
         device_rel_path, host_rel_path = t.split(':')
         host_abs_path = os.path.join(host_paths.DIR_SOURCE_ROOT, host_rel_path)
         self._data_deps.extend(
             [(host_abs_path,
               [None, 'chrome', 'test', 'data', device_rel_path])])
 
   def GetDataDependencies(self):
     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(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'])
-
-        # 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:
       for m in c['methods']:
         a = dict(c['annotations'])
         a.update(m['annotations'])
         inflated_tests.append({
             'class': c['class'],
             'method': m['method'],
             'annotations': a,
@@ skipping 36 matching lines @@
   @staticmethod
   def GenerateTestResults(
       result_code, result_bundle, statuses, start_ms, duration_ms):
     return GenerateTestResults(result_code, result_bundle, statuses,
                                start_ms, duration_ms)
 
   #override
   def TearDown(self):
     if self._isolate_delegate:
       self._isolate_delegate.Clear()
-