michael@0: # Copyright (c) 2012 The Chromium Authors. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: """Gathers information about APKs."""
michael@0: 
michael@0: import collections
michael@0: import os
michael@0: import re
michael@0: 
michael@0: import cmd_helper
michael@0: 
michael@0: 
michael@0: class ApkInfo(object):
michael@0:   """Helper class for inspecting APKs."""
michael@0:   _PROGUARD_PATH = os.path.join(os.environ['ANDROID_SDK_ROOT'],
michael@0:                                 'tools/proguard/bin/proguard.sh')
michael@0:   if not os.path.exists(_PROGUARD_PATH):
michael@0:     _PROGUARD_PATH = os.path.join(os.environ['ANDROID_BUILD_TOP'],
michael@0:                                   'external/proguard/bin/proguard.sh')
michael@0:   _PROGUARD_CLASS_RE = re.compile(r'\s*?- Program class:\s*([\S]+)$')
michael@0:   _PROGUARD_METHOD_RE = re.compile(r'\s*?- Method:\s*(\S*)[(].*$')
michael@0:   _PROGUARD_ANNOTATION_RE = re.compile(r'\s*?- Annotation \[L(\S*);\]:$')
michael@0:   _PROGUARD_ANNOTATION_CONST_RE = re.compile(r'\s*?- Constant element value.*$')
michael@0:   _PROGUARD_ANNOTATION_VALUE_RE = re.compile(r'\s*?- \S+? \[(.*)\]$')
michael@0:   _AAPT_PACKAGE_NAME_RE = re.compile(r'package: .*name=\'(\S*)\'')
michael@0: 
michael@0:   def __init__(self, apk_path, jar_path):
michael@0:     if not os.path.exists(apk_path):
michael@0:       raise Exception('%s not found, please build it' % apk_path)
michael@0:     self._apk_path = apk_path
michael@0:     if not os.path.exists(jar_path):
michael@0:       raise Exception('%s not found, please build it' % jar_path)
michael@0:     self._jar_path = jar_path
michael@0:     self._annotation_map = collections.defaultdict(list)
michael@0:     self._test_methods = []
michael@0:     self._Initialize()
michael@0: 
michael@0:   def _Initialize(self):
michael@0:     proguard_output = cmd_helper.GetCmdOutput([self._PROGUARD_PATH,
michael@0:                                                '-injars', self._jar_path,
michael@0:                                                '-dontshrink',
michael@0:                                                '-dontoptimize',
michael@0:                                                '-dontobfuscate',
michael@0:                                                '-dontpreverify',
michael@0:                                                '-dump',
michael@0:                                               ]).split('\n')
michael@0:     clazz = None
michael@0:     method = None
michael@0:     annotation = None
michael@0:     has_value = False
michael@0:     qualified_method = None
michael@0:     for line in proguard_output:
michael@0:       m = self._PROGUARD_CLASS_RE.match(line)
michael@0:       if m:
michael@0:         clazz = m.group(1).replace('/', '.')  # Change package delim.
michael@0:         annotation = None
michael@0:         continue
michael@0:       m = self._PROGUARD_METHOD_RE.match(line)
michael@0:       if m:
michael@0:         method = m.group(1)
michael@0:         annotation = None
michael@0:         qualified_method = clazz + '#' + method
michael@0:         if method.startswith('test') and clazz.endswith('Test'):
michael@0:           self._test_methods += [qualified_method]
michael@0:         continue
michael@0:       m = self._PROGUARD_ANNOTATION_RE.match(line)
michael@0:       if m:
michael@0:         assert qualified_method
michael@0:         annotation = m.group(1).split('/')[-1]  # Ignore the annotation package.
michael@0:         self._annotation_map[qualified_method].append(annotation)
michael@0:         has_value = False
michael@0:         continue
michael@0:       if annotation:
michael@0:         assert qualified_method
michael@0:         if not has_value:
michael@0:           m = self._PROGUARD_ANNOTATION_CONST_RE.match(line)
michael@0:           if m:
michael@0:             has_value = True
michael@0:         else:
michael@0:           m = self._PROGUARD_ANNOTATION_VALUE_RE.match(line)
michael@0:           if m:
michael@0:             value = m.group(1)
michael@0:             self._annotation_map[qualified_method].append(
michael@0:                 annotation + ':' + value)
michael@0:             has_value = False
michael@0: 
michael@0:   def _GetAnnotationMap(self):
michael@0:     return self._annotation_map
michael@0: 
michael@0:   def _IsTestMethod(self, test):
michael@0:     class_name, method = test.split('#')
michael@0:     return class_name.endswith('Test') and method.startswith('test')
michael@0: 
michael@0:   def GetApkPath(self):
michael@0:     return self._apk_path
michael@0: 
michael@0:   def GetPackageName(self):
michael@0:     """Returns the package name of this APK."""
michael@0:     aapt_output = cmd_helper.GetCmdOutput(
michael@0:         ['aapt', 'dump', 'badging', self._apk_path]).split('\n')
michael@0:     for line in aapt_output:
michael@0:       m = self._AAPT_PACKAGE_NAME_RE.match(line)
michael@0:       if m:
michael@0:         return m.group(1)
michael@0:     raise Exception('Failed to determine package name of %s' % self._apk_path)
michael@0: 
michael@0:   def GetTestAnnotations(self, test):
michael@0:     """Returns a list of all annotations for the given |test|. May be empty."""
michael@0:     if not self._IsTestMethod(test):
michael@0:       return []
michael@0:     return self._GetAnnotationMap()[test]
michael@0: 
michael@0:   def _AnnotationsMatchFilters(self, annotation_filter_list, annotations):
michael@0:     """Checks if annotations match any of the filters."""
michael@0:     if not annotation_filter_list:
michael@0:       return True
michael@0:     for annotation_filter in annotation_filter_list:
michael@0:       filters = annotation_filter.split('=')
michael@0:       if len(filters) == 2:
michael@0:         key = filters[0]
michael@0:         value_list = filters[1].split(',')
michael@0:         for value in value_list:
michael@0:           if key + ':' + value in annotations:
michael@0:             return True
michael@0:       elif annotation_filter in annotations:
michael@0:         return True
michael@0:     return False
michael@0: 
michael@0:   def GetAnnotatedTests(self, annotation_filter_list):
michael@0:     """Returns a list of all tests that match the given annotation filters."""
michael@0:     return [test for test, annotations in self._GetAnnotationMap().iteritems()
michael@0:             if self._IsTestMethod(test) and self._AnnotationsMatchFilters(
michael@0:                 annotation_filter_list, annotations)]
michael@0: 
michael@0:   def GetTestMethods(self):
michael@0:     """Returns a list of all test methods in this apk as Class#testMethod."""
michael@0:     return self._test_methods
michael@0: 
michael@0:   @staticmethod
michael@0:   def IsPythonDrivenTest(test):
michael@0:     return 'pythonDrivenTests' in test