1# Copyright 2018, The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Integration Finder class."""
16
17import copy
18import logging
19import os
20import re
21import tempfile
22import xml.etree.ElementTree as ElementTree
23from zipfile import ZipFile
24
25from atest import atest_error
26from atest import atest_utils
27from atest import constants
28from atest.test_finders import test_filter_utils
29from atest.test_finders import test_finder_base
30from atest.test_finders import test_finder_utils
31from atest.test_finders import test_info
32from atest.test_runners import atest_tf_test_runner
33
34# Find integration name based on file path of integration config xml file.
35# Group matches "foo/bar" given "blah/res/config/foo/bar.xml from source code
36# res directory or "blah/config/foo/bar.xml from prebuilt jars.
37_INT_NAME_RE = re.compile(r'^.*\/config\/(?P<int_name>.*).xml$')
38_TF_TARGETS = frozenset(['tradefed', 'tradefed-contrib'])
39_GTF_TARGETS = frozenset(['google-tradefed', 'google-tradefed-contrib'])
40_CONTRIB_TARGETS = frozenset(['google-tradefed-contrib'])
41_TF_RES_DIRS = frozenset(['../res/config', 'res/config'])
42
43
44class TFIntegrationFinder(test_finder_base.TestFinderBase):
45  """Integration Finder class."""
46
47  NAME = 'INTEGRATION'
48  _TEST_RUNNER = atest_tf_test_runner.AtestTradefedTestRunner.NAME
49
50  def __init__(self, module_info=None):
51    super().__init__()
52    self.root_dir = os.environ.get(constants.ANDROID_BUILD_TOP)
53    self.module_info = module_info
54    # TODO: Break this up into AOSP/google_tf integration finders.
55    self.tf_dirs, self.gtf_dirs = self._get_integration_dirs()
56    self.integration_dirs = self.tf_dirs + self.gtf_dirs
57    self.temp_dir = tempfile.TemporaryDirectory()
58
59  def _get_mod_paths(self, module_name):
60    """Return the paths of the given module name."""
61    if self.module_info:
62      # Since aosp/801774 merged, the path of test configs have been
63      # changed to ../res/config.
64      if module_name in _CONTRIB_TARGETS:
65        mod_paths = self.module_info.get_paths(module_name)
66        return [
67            os.path.join(path, res_path)
68            for path in mod_paths
69            for res_path in _TF_RES_DIRS
70        ]
71      return self.module_info.get_paths(module_name)
72    return []
73
74  def _get_integration_dirs(self):
75    """Get integration dirs from MODULE_INFO based on targets.
76
77    Returns:
78        A tuple of lists of strings of integration dir rel to repo root.
79    """
80    tf_dirs = list(
81        filter(None, [d for x in _TF_TARGETS for d in self._get_mod_paths(x)])
82    )
83    gtf_dirs = list(
84        filter(None, [d for x in _GTF_TARGETS for d in self._get_mod_paths(x)])
85    )
86    return tf_dirs, gtf_dirs
87
88  def _get_build_targets(self, rel_config):
89    config_file = os.path.join(self.root_dir, rel_config)
90    xml_root = self._load_xml_file(config_file)
91    targets = test_finder_utils.get_targets_from_xml_root(
92        xml_root, self.module_info
93    )
94    if self.gtf_dirs:
95      targets.add(constants.GTF_TARGET)
96    return frozenset(targets)
97
98  def _load_xml_file(self, path):
99    """Load an xml file with option to expand <include> tags
100
101    Args:
102        path: A string of path to xml file.
103
104    Returns:
105        An xml.etree.ElementTree.Element instance of the root of the tree.
106    """
107    tree = ElementTree.parse(path)
108    root = tree.getroot()
109    self._load_include_tags(root)
110    return root
111
112  # pylint: disable=invalid-name
113  def _load_include_tags(self, root):
114    """Recursively expand in-place the <include> tags in a given xml tree.
115
116    Python xml libraries don't support our type of <include> tags. Logic
117    used below is modified version of the built-in ElementInclude logic
118    found here:
119    https://github.com/python/cpython/blob/2.7/Lib/xml/etree/ElementInclude.py
120
121    Args:
122        root: The root xml.etree.ElementTree.Element.
123
124    Returns:
125        An xml.etree.ElementTree.Element instance with
126        include tags expanded.
127    """
128    i = 0
129    while i < len(root):
130      elem = root[i]
131      if elem.tag == 'include':
132        # expand included xml file
133        integration_name = elem.get('name')
134        if not integration_name:
135          atest_utils.print_and_log_warning(
136              'skipping <include> tag with no "name" value'
137          )
138          continue
139        full_paths = self._search_integration_dirs(integration_name)
140        if not full_paths:
141          full_paths = self._search_prebuilt_jars(integration_name)
142        node = None
143        if full_paths:
144          node = self._load_xml_file(full_paths[0])
145        if node is None:
146          raise atest_error.FatalIncludeError(
147              "can't load %r" % integration_name
148          )
149        node = copy.copy(node)
150        if elem.tail:
151          node.tail = (node.tail or '') + elem.tail
152        root[i] = node
153      i = i + 1
154
155  def _search_integration_dirs(self, name):
156    """Search integration dirs for name and return full path.
157
158    Args:
159        name: A string of integration name as seen in tf's list configs.
160
161    Returns:
162        A list of test path.
163    """
164    test_files = []
165    for integration_dir in self.integration_dirs:
166      abs_path = os.path.join(self.root_dir, integration_dir)
167      found_test_files = test_finder_utils.run_find_cmd(
168          test_finder_utils.TestReferenceType.INTEGRATION, abs_path, name
169      )
170      if found_test_files:
171        test_files.extend(found_test_files)
172    return test_files
173
174  def find_test_by_integration_name(self, name):
175    """Find the test info matching the given integration name.
176
177    Args:
178        name: A string of integration name as seen in tf's list configs.
179
180    Returns:
181        A populated TestInfo namedtuple if test found, else None
182    """
183    class_name = None
184    parse_result = test_finder_utils.parse_test_reference(name)
185    if parse_result:
186      name = parse_result['module_name']
187      class_name = parse_result['pkg_class_name']
188      method = parse_result.get('method_name', '')
189      if method:
190        class_name = class_name + '#' + method
191    test_files = self._search_integration_dirs(name)
192    if not test_files:
193      # Check prebuilt jars if input name is in jars.
194      test_files = self._search_prebuilt_jars(name)
195    # Don't use names that simply match the path,
196    # must be the actual name used by TF to run the test.
197    t_infos = []
198    for test_file in test_files:
199      t_info = self._get_test_info(name, test_file, class_name)
200      if t_info:
201        t_infos.append(t_info)
202    return t_infos
203
204  def _get_prebuilt_jars(self):
205    """Get prebuilt jars based on targets.
206
207    Returns:
208        A tuple of lists of strings of prebuilt jars.
209    """
210    prebuilt_jars = []
211    for tf_dir in self.tf_dirs:
212      for tf_target in _TF_TARGETS:
213        jar_path = os.path.join(
214            self.root_dir,
215            tf_dir,
216            '..',
217            'filegroups',
218            'tradefed',
219            tf_target + '.jar',
220        )
221        if os.path.exists(jar_path):
222          prebuilt_jars.append(jar_path)
223    for gtf_dir in self.gtf_dirs:
224      for gtf_target in _GTF_TARGETS:
225        jar_path = os.path.join(
226            self.root_dir,
227            gtf_dir,
228            '..',
229            'filegroups',
230            'google-tradefed',
231            gtf_target + '.jar',
232        )
233        if os.path.exists(jar_path):
234          prebuilt_jars.append(jar_path)
235    return prebuilt_jars
236
237  def _search_prebuilt_jars(self, name):
238    """Search tradefed prebuilt jar which has matched name.
239
240    Search if input name matched prebuilt tradefed jar. If matched, extract
241    the jar file to temp directly for later on test info handling.
242
243    Args:
244        name: A string of integration name as seen in tf's list configs.
245
246    Returns:
247        A list of test path.
248    """
249
250    xml_path = 'config/{}.xml'.format(name)
251    test_files = []
252    prebuilt_jars = self._get_prebuilt_jars()
253    logging.debug('Found prebuilt_jars=%s', prebuilt_jars)
254    for prebuilt_jar in prebuilt_jars:
255      with ZipFile(prebuilt_jar, 'r') as jar_file:
256        jar_contents = jar_file.namelist()
257        if xml_path in jar_contents:
258          extract_path = os.path.join(
259              self.temp_dir.name, os.path.basename(prebuilt_jar)
260          )
261          if not os.path.exists(extract_path):
262            logging.debug('Extracting %s to %s', prebuilt_jar, extract_path)
263            jar_file.extractall(extract_path)
264          test_files.append(os.path.join(extract_path, xml_path))
265    return test_files
266
267  def _get_test_info(self, name, test_file, class_name):
268    """Find the test info matching the given test_file and class_name.
269
270    Args:
271        name: A string of integration name as seen in tf's list configs.
272        test_file: A string of test_file full path.
273        class_name: A string of user's input.
274
275    Returns:
276        A populated TestInfo namedtuple if test found, else None.
277    """
278    match = _INT_NAME_RE.match(test_file)
279    if not match:
280      atest_utils.print_and_log_error(
281          'Integration test outside config dir: %s', test_file
282      )
283      return None
284    int_name = match.group('int_name')
285    if int_name != name:
286      logging.debug(
287          'Input (%s) not valid integration name, did you mean: %s?',
288          name,
289          int_name,
290      )
291      return None
292    rel_config = os.path.relpath(test_file, self.root_dir)
293    filters = frozenset()
294    if class_name:
295      class_name, methods = test_filter_utils.split_methods(class_name)
296      test_filters = []
297      if '.' in class_name:
298        test_filters.append(test_info.TestFilter(class_name, methods))
299      else:
300        logging.debug(
301            'Looking up fully qualified class name for: %s.'
302            'Improve speed by using fully qualified names.',
303            class_name,
304        )
305        paths = test_finder_utils.find_class_file(self.root_dir, class_name)
306        if not paths:
307          return None
308        for path in paths:
309          class_name = test_filter_utils.get_fully_qualified_class_name(path)
310          test_filters.append(test_info.TestFilter(class_name, methods))
311      filters = frozenset(test_filters)
312    return test_info.TestInfo(
313        test_name=name,
314        test_runner=self._TEST_RUNNER,
315        build_targets=self._get_build_targets(rel_config),
316        data={
317            constants.TI_REL_CONFIG: rel_config,
318            constants.TI_FILTER: filters,
319        },
320    )
321
322  def find_int_test_by_path(self, path):
323    """Find the first test info matching the given path.
324
325    Strategy:
326        path_to_integration_file --> Resolve to INTEGRATION
327        # If the path is a dir, we return nothing.
328        path_to_dir_with_integration_files --> Return None
329
330    Args:
331        path: A string of the test's path.
332
333    Returns:
334        A list of populated TestInfo namedtuple if test found, else None
335    """
336    path, _ = test_filter_utils.split_methods(path)
337
338    # Make sure we're looking for a config.
339    if not path.endswith('.xml'):
340      return None
341
342    # TODO: See if this can be generalized and shared with methods above
343    # create absolute path from cwd and remove symbolic links
344    path = os.path.realpath(path)
345    if not os.path.exists(path):
346      logging.debug('"%s": file not found!', path)
347      return None
348    int_dir = test_finder_utils.get_int_dir_from_path(
349        path, self.integration_dirs
350    )
351    if int_dir:
352      rel_config = os.path.relpath(path, self.root_dir)
353      match = _INT_NAME_RE.match(rel_config)
354      if not match:
355        atest_utils.print_and_log_error(
356            'Integration test outside config dir: %s', rel_config
357        )
358        return None
359      int_name = match.group('int_name')
360      return [
361          test_info.TestInfo(
362              test_name=int_name,
363              test_runner=self._TEST_RUNNER,
364              build_targets=self._get_build_targets(rel_config),
365              data={
366                  constants.TI_REL_CONFIG: rel_config,
367                  constants.TI_FILTER: frozenset(),
368              },
369          )
370      ]
371    return None
372