1# Copyright 2017, 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"""Atest Tradefed test runner class."""
16
17# pylint: disable=too-many-lines
18
19from __future__ import annotations
20from __future__ import print_function
21
22from abc import ABC, abstractmethod
23import dataclasses
24import enum
25from functools import partial
26import json
27import logging
28import os
29from pathlib import Path
30import re
31import select
32import shutil
33import socket
34import time
35from typing import Any, Dict, List, Set, Tuple
36
37from atest import atest_configs
38from atest import atest_error
39from atest import atest_utils
40from atest import constants
41from atest import module_info
42from atest import result_reporter
43from atest.atest_enum import DetectType, ExitCode
44from atest.coverage import coverage
45from atest.logstorage import logstorage_utils
46from atest.metrics import metrics
47from atest.test_finders import test_finder_utils
48from atest.test_finders import test_info
49from atest.test_runner_invocation import TestRunnerInvocation
50from atest.test_runners import test_runner_base as trb
51from atest.test_runners.event_handler import EventHandler
52
53POLL_FREQ_SECS = 10
54SOCKET_HOST = '127.0.0.1'
55SOCKET_QUEUE_MAX = 1
56SOCKET_BUFFER = 4096
57SELECT_TIMEOUT = 0.5
58
59# Socket Events of form FIRST_EVENT {JSON_DATA}\nSECOND_EVENT {JSON_DATA}
60# EVENT_RE has groups for the name and the data. "." does not match \n.
61EVENT_RE = re.compile(
62    r'\n*(?P<event_name>[A-Z_]+) (?P<json_data>{.*})(?=\n|.)*'
63)
64
65# Remove aapt from build dependency, use prebuilt version instead.
66EXEC_DEPENDENCIES = ('adb', 'fastboot')
67
68LOG_FOLDER_NAME = 'log'
69
70_INTEGRATION_FINDERS = frozenset(['', 'INTEGRATION', 'INTEGRATION_FILE_PATH'])
71
72# AAPT binary name
73_AAPT = 'aapt'
74
75# The exist code mapping of tradefed.
76_TF_EXIT_CODE = [
77    'NO_ERROR',
78    'CONFIG_EXCEPTION',
79    'NO_BUILD',
80    'DEVICE_UNRESPONSIVE',
81    'DEVICE_UNAVAILABLE',
82    'FATAL_HOST_ERROR',
83    'THROWABLE_EXCEPTION',
84    'NO_DEVICE_ALLOCATED',
85    'WRONG_JAVA_VERSION',
86]
87
88
89class Error(Exception):
90  """Module-level error."""
91
92
93class TradeFedExitError(Error):
94  """Raised when TradeFed exists before test run has finished."""
95
96  def __init__(self, exit_code):
97    super().__init__()
98    self.exit_code = exit_code
99
100  def __str__(self):
101    tf_error_reason = self._get_exit_reason(self.exit_code)
102    return (
103        'TradeFed subprocess exited early with exit code='
104        f'{self.exit_code}({tf_error_reason}).'
105    )
106
107  def _get_exit_reason(self, exit_code):
108    if 0 < exit_code < len(_TF_EXIT_CODE):
109      return atest_utils.mark_red(_TF_EXIT_CODE[exit_code])
110    return 'Unknown exit status'
111
112
113class AtestTradefedTestRunner(trb.TestRunnerBase):
114  """TradeFed Test Runner class."""
115
116  NAME = 'AtestTradefedTestRunner'
117  EXECUTABLE = 'atest_tradefed.sh'
118  # Common base template used by all TF tests
119  _TF_LOCAL_MIN = 'template/atest_local_min'
120  # Base template used by the device tests (tests requires device to run)
121  _TF_DEVICE_TEST_TEMPLATE = 'template/atest_device_test_base'
122  # Base template used by the deviceless tests
123  _TF_DEVICELESS_TEST_TEMPLATE = 'template/atest_deviceless_test_base'
124  # Use --no-enable-granular-attempts to control reporter replay behavior.
125  # TODO(b/142630648): Enable option enable-granular-attempts
126  # in sharding mode.
127  _LOG_ARGS = (
128      '--{log_root_option_name}={log_path} '
129      '{log_ext_option} '
130      '--no-enable-granular-attempts'
131  )
132  _RUN_CMD = (
133      '{env} {exe} {template} '
134      '--template:map test=atest '
135      '--template:map log_saver={log_saver} '
136      '{tf_customize_template} {log_args} {args}'
137  )
138  _BUILD_REQ = {'tradefed-core'}
139  _RERUN_OPTION_GROUP = [
140      constants.ITERATIONS,
141      constants.RERUN_UNTIL_FAILURE,
142      constants.RETRY_ANY_FAILURE,
143  ]
144
145  # We're using a class attribute because we're recreating runner instances
146  # for different purposes throughout an invocation.
147  # TODO(b/283352341): Remove this once we refactor to have runner instances.
148  _MINIMAL_BUILD_TARGETS = set()
149
150  def __init__(
151      self,
152      results_dir: str,
153      extra_args: Dict[str, Any],
154      mod_info: module_info.ModuleInfo = None,
155      minimal_build: bool = False,
156      **kwargs,
157  ):
158    """Init stuff for base class."""
159    super().__init__(results_dir, **kwargs)
160    self.module_info = mod_info
161    self.log_path = os.path.join(results_dir, LOG_FOLDER_NAME)
162    # (b/275537997) results_dir could be '' in test_runner_handler; only
163    # mkdir when it is invoked by run_tests.
164    if results_dir:
165      Path(self.log_path).mkdir(parents=True, exist_ok=True)
166    self.log_args = {
167        'log_root_option_name': constants.LOG_ROOT_OPTION_NAME,
168        'log_ext_option': constants.LOG_SAVER_EXT_OPTION,
169        'log_path': self.log_path,
170        'proto_path': os.path.join(
171            self.results_dir, constants.ATEST_TEST_RECORD_PROTO
172        ),
173    }
174    self.run_cmd_dict = {
175        'env': '',
176        'exe': self.EXECUTABLE,
177        'template': self._TF_LOCAL_MIN,
178        'log_saver': constants.ATEST_TF_LOG_SAVER,
179        'tf_customize_template': '',
180        'args': '',
181        'log_args': self._LOG_ARGS.format(**self.log_args),
182    }
183    self.root_dir = os.environ.get(constants.ANDROID_BUILD_TOP)
184    self._is_host_enabled = extra_args.get(constants.HOST, False)
185    self._minimal_build = minimal_build
186    logging.debug('Enable minimal build: %s' % self._minimal_build)
187    metrics.LocalDetectEvent(
188        detect_type=DetectType.IS_MINIMAL_BUILD, result=int(self._minimal_build)
189    )
190
191  def requires_device_update(
192      self, test_infos: List[test_info.TestInfo]
193  ) -> bool:
194    """Checks whether this runner requires device update."""
195
196    requires_device_update = False
197    for info in test_infos:
198      test = self._create_test(info)
199      requires_device_update |= test.requires_device_update()
200
201    return requires_device_update
202
203  # @typing.override
204  def create_invocations(
205      self,
206      extra_args: Dict[str, Any],
207      test_infos: List[test_info.TestInfo],
208  ) -> List[TestRunnerInvocation]:
209    """Create separate test runner invocations for device and deviceless tests.
210
211    Args:
212        extra_args: Dict of extra args to pass to the invocations
213        test_infos: A list of TestInfos.
214
215    Returns:
216        A list of TestRunnerInvocation instances.
217    """
218    invocations = []
219    device_test_infos, deviceless_test_infos = self._partition_tests(test_infos)
220    if deviceless_test_infos:
221      extra_args_for_deviceless_test = extra_args.copy()
222      extra_args_for_deviceless_test.update({constants.HOST: True})
223      invocations.append(
224          TestRunnerInvocation(
225              test_runner=self,
226              extra_args=extra_args_for_deviceless_test,
227              test_infos=deviceless_test_infos,
228          )
229      )
230    if device_test_infos:
231      invocations.append(
232          TestRunnerInvocation(
233              test_runner=self,
234              extra_args=extra_args,
235              test_infos=device_test_infos,
236          )
237      )
238
239    return invocations
240
241  def _partition_tests(
242      self,
243      test_infos: List[test_info.TestInfo],
244  ) -> (List[test_info.TestInfo], List[test_info.TestInfo]):
245    """Partition input tests into two lists based on whether it requires device.
246
247    Args:
248        test_infos: A list of TestInfos.
249
250    Returns:
251        Two lists one contains device tests the other contains deviceless tests.
252    """
253    device_test_infos = []
254    deviceless_test_infos = []
255
256    for info in test_infos:
257      test = self._create_test(info)
258      if test.requires_device():
259        device_test_infos.append(info)
260      else:
261        deviceless_test_infos.append(info)
262
263    return device_test_infos, deviceless_test_infos
264
265  def _try_set_gts_authentication_key(self):
266    """Set GTS authentication key if it is available or exists.
267
268    Strategy:
269        Get APE_API_KEY from os.environ:
270            - If APE_API_KEY is already set by user -> do nothing.
271        Get the APE_API_KEY from constants:
272            - If the key file exists -> set to env var.
273        If APE_API_KEY isn't set and the key file doesn't exist:
274            - Warn user some GTS tests may fail without authentication.
275    """
276    if os.environ.get('APE_API_KEY'):
277      logging.debug('APE_API_KEY is set by developer.')
278      return
279    ape_api_key = constants.GTS_GOOGLE_SERVICE_ACCOUNT
280    key_path = os.path.join(self.root_dir, ape_api_key)
281    if ape_api_key and os.path.exists(key_path):
282      logging.debug('Set APE_API_KEY: %s', ape_api_key)
283      os.environ['APE_API_KEY'] = key_path
284    else:
285      logging.debug(
286          'APE_API_KEY not set, some GTS tests may fail without authentication.'
287      )
288
289  def run_tests(self, test_infos, extra_args, reporter):
290    """Run the list of test_infos. See base class for more.
291
292    Args:
293        test_infos: A list of TestInfos.
294        extra_args: Dict of extra args to add to test run.
295        reporter: An instance of result_report.ResultReporter.
296
297    Returns:
298        0 if tests succeed, non-zero otherwise.
299    """
300    logging.debug('TF test runner running tests %s', test_infos)
301    reporter.log_path = self.log_path
302    reporter.rerun_options = self._extract_rerun_options(extra_args)
303    # Set google service key if it's available or found before
304    # running tests.
305    self._try_set_gts_authentication_key()
306    result = 0
307    upload_start = time.time()
308    creds, inv = (
309        logstorage_utils.do_upload_flow(extra_args)
310        if logstorage_utils.is_upload_enabled(extra_args)
311        else (None, None)
312    )
313    metrics.LocalDetectEvent(
314        detect_type=DetectType.UPLOAD_FLOW_MS,
315        result=int((time.time() - upload_start) * 1000),
316    )
317    try:
318      verify_key = atest_utils.get_verify_key(
319          [test_infos[0].test_name], extra_args
320      )
321      # Change CWD to repo root to ensure TF can find prebuilt SDKs
322      # for some path-sensitive tests like robolectric.
323      os.chdir(os.path.abspath(os.getenv(constants.ANDROID_BUILD_TOP)))
324
325      # Copy symbols if there are tests belong to native test.
326      self._handle_native_tests(test_infos)
327
328      if os.getenv(trb.OLD_OUTPUT_ENV_VAR):
329        result = self.run_tests_raw(test_infos, extra_args, reporter)
330      else:
331        result = self.run_tests_pretty(test_infos, extra_args, reporter)
332    except atest_error.DryRunVerificationError as e:
333      atest_utils.colorful_print(str(e), constants.RED)
334      return ExitCode.VERIFY_FAILURE
335    finally:
336      if inv:
337        try:
338          logging.disable(logging.INFO)
339          # Always set invocation status to completed due to the ATest
340          # handle whole process by its own.
341          inv['schedulerState'] = 'completed'
342          logstorage_utils.BuildClient(creds).update_invocation(inv)
343          reporter.test_result_link = (
344              constants.RESULT_LINK % inv['invocationId']
345          )
346        finally:
347          logging.disable(logging.NOTSET)
348    return result
349
350  def run_tests_raw(self, test_infos, extra_args, reporter):
351    """Run the list of test_infos. See base class for more.
352
353    Args:
354        test_infos: A list of TestInfos.
355        extra_args: Dict of extra args to add to test run.
356        reporter: An instance of result_report.ResultReporter.
357
358    Returns:
359        0 if tests succeed, non-zero otherwise.
360    """
361    reporter.register_unsupported_runner(self.NAME)
362
363    ret_code = ExitCode.SUCCESS
364    run_cmds = self.generate_run_commands(test_infos, extra_args)
365    logging.debug('Running test: %s', run_cmds[0])
366    subproc = self.run(
367        run_cmds[0],
368        output_to_stdout=True,
369        env_vars=self.generate_env_vars(extra_args),
370    )
371    ret_code |= self.wait_for_subprocess(subproc)
372    return ret_code
373
374  def run_tests_pretty(self, test_infos, extra_args, reporter):
375    """Run the list of test_infos. See base class for more.
376
377    Args:
378        test_infos: A list of TestInfos.
379        extra_args: Dict of extra args to add to test run.
380        reporter: An instance of result_report.ResultReporter.
381
382    Returns:
383        0 if tests succeed, non-zero otherwise.
384    """
385    ret_code = ExitCode.SUCCESS
386    server = self._start_socket_server()
387    run_cmds = self.generate_run_commands(
388        test_infos, extra_args, server.getsockname()[1]
389    )
390    logging.debug('Running test: %s', run_cmds[0])
391    subproc = self.run(
392        run_cmds[0],
393        output_to_stdout=extra_args.get(constants.VERBOSE, False),
394        env_vars=self.generate_env_vars(extra_args),
395    )
396    self.handle_subprocess(
397        subproc,
398        partial(self._start_monitor, server, subproc, reporter, extra_args),
399    )
400    server.close()
401    ret_code |= self.wait_for_subprocess(subproc)
402    return ret_code
403
404  # pylint: disable=too-many-branches
405  # pylint: disable=too-many-locals
406  def _start_monitor(self, server, tf_subproc, reporter, extra_args):
407    """Polling and process event.
408
409    Args:
410        server: Socket server object.
411        tf_subproc: The tradefed subprocess to poll.
412        reporter: Result_Reporter object.
413        extra_args: Dict of extra args to add to test run.
414    """
415    inputs = [server]
416    event_handlers = {}
417    data_map = {}
418    inv_socket = None
419    while inputs:
420      try:
421        readable, _, _ = select.select(inputs, [], [], SELECT_TIMEOUT)
422        for socket_object in readable:
423          if socket_object is server:
424            conn, addr = socket_object.accept()
425            logging.debug('Accepted connection from %s', addr)
426            conn.setblocking(False)
427            inputs.append(conn)
428            data_map[conn] = ''
429            # The First connection should be invocation
430            # level reporter.
431            if not inv_socket:
432              inv_socket = conn
433          else:
434            # Count invocation level reporter events
435            # without showing real-time information.
436            if inv_socket == socket_object:
437              reporter.silent = True
438              event_handler = event_handlers.setdefault(
439                  socket_object, EventHandler(reporter, self.NAME)
440              )
441            else:
442              event_handler = event_handlers.setdefault(
443                  socket_object,
444                  EventHandler(
445                      result_reporter.ResultReporter(
446                          collect_only=extra_args.get(
447                              constants.COLLECT_TESTS_ONLY
448                          ),
449                      ),
450                      self.NAME,
451                  ),
452              )
453            recv_data = self._process_connection(
454                data_map, socket_object, event_handler
455            )
456            if not recv_data:
457              inputs.remove(socket_object)
458              socket_object.close()
459      finally:
460        # Subprocess ended and all socket clients were closed.
461        if tf_subproc.poll() is not None and len(inputs) == 1:
462          inputs.pop().close()
463          if not reporter.all_test_results:
464            if atest_configs.GLOBAL_ARGS.user_type:
465              atest_utils.colorful_print(
466                  "The test module doesn't support "
467                  f"'{atest_configs.GLOBAL_ARGS.user_type}' "
468                  'user type, please check test config.',
469                  constants.RED,
470              )
471            atest_utils.colorful_print(
472                r'No test results available. TradeFed did not find'
473                r' any test cases to run. This is possibly due to'
474                r' the no tests matching the current test filters'
475                r' or misconfigured AndroidTest.xml. Test Logs'
476                r' are saved in '
477                f'{reporter.log_path}.',
478                constants.RED,
479                constants.WHITE,
480            )
481          if not data_map:
482            metrics.LocalDetectEvent(
483                detect_type=DetectType.TF_EXIT_CODE,
484                result=tf_subproc.returncode,
485            )
486            raise TradeFedExitError(tf_subproc.returncode)
487          self._handle_log_associations(event_handlers)
488
489  def _process_connection(self, data_map, conn, event_handler):
490    """Process a socket connection between TF and ATest.
491
492    Expect data of form EVENT_NAME {JSON_DATA}.  Multiple events will be
493    \n deliminated.  Need to buffer data in case data exceeds socket
494    buffer.
495    E.q.
496        TEST_RUN_STARTED {runName":"hello_world_test","runAttempt":0}\n
497        TEST_STARTED {"start_time":2172917, "testName":"PrintHelloWorld"}\n
498    Args:
499        data_map: The data map of all connections.
500        conn: Socket connection.
501        event_handler: EventHandler object.
502
503    Returns:
504        True if conn.recv() has data , False otherwise.
505    """
506    # Set connection into blocking mode.
507    conn.settimeout(None)
508    data = conn.recv(SOCKET_BUFFER)
509    if isinstance(data, bytes):
510      data = data.decode()
511    logging.debug('received: %s', data)
512    if data:
513      data_map[conn] += data
514      while True:
515        match = EVENT_RE.match(data_map[conn])
516        if not match:
517          break
518        try:
519          event_data = json.loads(match.group('json_data'))
520        except ValueError:
521          logging.debug('Json incomplete, wait for more data')
522          break
523        event_name = match.group('event_name')
524        event_handler.process_event(event_name, event_data)
525        data_map[conn] = data_map[conn][match.end() :]
526    return bool(data)
527
528  def _start_socket_server(self):
529    """Start a TCP server."""
530    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
531    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
532    # Port 0 lets the OS pick an open port between 1024 and 65535.
533    server.bind((SOCKET_HOST, 0))
534    server.listen(SOCKET_QUEUE_MAX)
535    server.settimeout(POLL_FREQ_SECS)
536    logging.debug('Socket server started on port %s', server.getsockname()[1])
537    return server
538
539  def generate_env_vars(self, extra_args):
540    """Convert extra args and test_infos into env vars.
541
542    Args:
543        extra_args: Dict of extra args to add to test run.
544        test_infos: A list of TestInfos.
545
546    Returns:
547        A dict modified from os.getenv.copy().
548    """
549    env_vars = os.environ.copy()
550    if constants.TF_GLOBAL_CONFIG and is_log_upload_enabled(extra_args):
551      env_vars['TF_GLOBAL_CONFIG'] = constants.TF_GLOBAL_CONFIG
552    debug_port = extra_args.get(constants.TF_DEBUG, '')
553    if debug_port:
554      env_vars['TF_DEBUG'] = 'true'
555      env_vars['TF_DEBUG_PORT'] = str(debug_port)
556
557    filtered_paths = []
558    for path in str(env_vars.get('PYTHONPATH', '')).split(':'):
559      # TODO (b/166216843) Remove the hacky PYTHON path workaround.
560      if (
561          str(path).startswith('/tmp/Soong.python_')
562          and str(path).find('googleapiclient') > 0
563      ):
564        continue
565      filtered_paths.append(path)
566    if filtered_paths:
567      env_vars['PYTHONPATH'] = ':'.join(filtered_paths)
568
569    # Use prebuilt aapt if there's no aapt under android system path which
570    # is aligned with build system.
571    # https://android.googlesource.com/platform/build/+/master/core/config.mk#529
572    if self._is_missing_exec(_AAPT):
573      prebuilt_aapt = Path.joinpath(
574          atest_utils.get_prebuilt_sdk_tools_dir(), _AAPT
575      )
576      if os.path.exists(prebuilt_aapt):
577        env_vars['PATH'] = str(prebuilt_aapt.parent) + ':' + env_vars['PATH']
578
579    # Add an env variable for the classpath that only contains the host jars
580    # required for the tests we'll be running.
581    if self._minimal_build:
582      self._generate_host_jars_env_var(env_vars)
583
584    return env_vars
585
586  def _generate_host_jars_env_var(self, env_vars):
587    def is_host_jar(p):
588      return p.suffix == '.jar' and p.is_relative_to(
589          Path(os.getenv(constants.ANDROID_HOST_OUT))
590      )
591
592    all_host_jars = []
593
594    for target in AtestTradefedTestRunner._MINIMAL_BUILD_TARGETS:
595      if target.variant != Variant.HOST:
596        continue
597      # Only use the first host jar because the same jar may be installed
598      # to multiple places.
599      module_host_jars = [
600          p
601          for p in self.module_info.get_installed_paths(target.module_name)
602          if is_host_jar(p)
603      ]
604      all_host_jars.extend(
605          [str(module_host_jars[0])] if module_host_jars else []
606      )
607
608    env_vars['ATEST_HOST_JARS'] = ':'.join(set(all_host_jars))
609    logging.debug(
610        'Set env ATEST_HOST_JARS: %s.', env_vars.get('ATEST_HOST_JARS')
611    )
612
613  # pylint: disable=unnecessary-pass
614  # Please keep above disable flag to ensure host_env_check is overridden.
615  def host_env_check(self):
616    """Check that host env has everything we need.
617
618    We actually can assume the host env is fine because we have the same
619    requirements that atest has. Update this to check for android env vars
620    if that changes.
621    """
622    pass
623
624  @staticmethod
625  def _is_missing_exec(executable):
626    """Check if system build executable is available.
627
628    Args:
629        executable: Executable we are checking for.
630
631    Returns:
632        True if executable is missing, False otherwise.
633    """
634    output = shutil.which(executable)
635    if not output:
636      return True
637    # TODO: Check if there is a clever way to determine if system adb is
638    # good enough.
639    root_dir = os.environ.get(constants.ANDROID_BUILD_TOP, '')
640    return os.path.commonprefix([output, root_dir]) != root_dir
641
642  def _use_minimal_build(self, test_infos: List[test_info.TestInfo]) -> bool:
643
644    if not self._minimal_build:
645      return False
646
647    unsupported = set()
648    for t_info in test_infos:
649      if t_info.test_finder in [
650          'CONFIG',
651          'INTEGRATION',
652          'INTEGRATION_FILE_PATH',
653      ]:
654        unsupported.add(t_info.test_name)
655      # For ltp and kselftest, keep it as no-minimal-build.
656      elif t_info.test_name in (
657          constants.REQUIRED_LTP_TEST_MODULES
658          + constants.REQUIRED_KSELFTEST_TEST_MODULES
659      ):
660        unsupported.add(t_info.test_name)
661
662    if not unsupported:
663      return True
664
665    atest_utils.print_and_log_warning(
666        'Minimal build was disabled because the following tests do not support'
667        ' it: %s',
668        unsupported,
669    )
670    return False
671
672  def get_test_runner_build_reqs(
673      self, test_infos: List[test_info.TestInfo]
674  ) -> Set[str]:
675    """Return the build requirements.
676
677    Args:
678        test_infos: List of TestInfo.
679
680    Returns:
681        Set of build targets.
682    """
683    if self._use_minimal_build(test_infos):
684      return self._get_test_runner_reqs_minimal(test_infos)
685
686    return self._get_test_runner_build_reqs_maximal(test_infos)
687
688  def _get_test_runner_build_reqs_maximal(
689      self, test_infos: List[test_info.TestInfo]
690  ) -> Set[str]:
691    build_req = self._BUILD_REQ.copy()
692    # Use different base build requirements if google-tf is around.
693    if self.module_info.is_module(constants.GTF_MODULE):
694      build_req = {constants.GTF_TARGET}
695    # Always add ATest's own TF target.
696    build_req.add(constants.ATEST_TF_MODULE)
697    # Add adb if we can't find it.
698    for executable in EXEC_DEPENDENCIES:
699      if self._is_missing_exec(executable):
700        if self.module_info.is_module(executable):
701          build_req.add(executable)
702
703    # Force rebuilt all jars under $ANDROID_HOST_OUT to prevent old version
704    # host jars break the test.
705    build_req |= self._get_host_framework_targets()
706
707    build_req |= trb.gather_build_targets(test_infos)
708    return build_req
709
710  def _get_test_runner_reqs_minimal(
711      self, test_infos: List[test_info.TestInfo]
712  ) -> Set[str]:
713
714    build_targets = set()
715    runtime_targets = set()
716
717    for info in test_infos:
718      test = self._create_test(info)
719      build_targets.update(test.query_build_targets())
720      runtime_targets.update(test.query_runtime_targets())
721
722    AtestTradefedTestRunner._MINIMAL_BUILD_TARGETS = runtime_targets
723
724    build_targets = {t.name() for t in build_targets}
725
726    return build_targets
727
728  def _create_test(self, t_info: test_info.TestInfo) -> Test:
729
730    info = self.module_info.get_module_info(t_info.raw_test_name)
731
732    if not info:
733      # In cases module info does not exist (e.g. TF integration tests), use the
734      # TestInfo to determine the test type. In the future we should ensure all
735      # tests have their corresponding module info and only rely on the module
736      # info to determine the test type.
737      atest_utils.print_and_log_warning(
738          'Could not find module information for %s', t_info.raw_test_name
739      )
740      return self._guess_test_type_for_missing_module(t_info)
741
742    def _select_variant(info):
743      variants = self.module_info.build_variants(info)
744      if len(variants) < 2:
745        return Variant.HOST if variants[0] == 'HOST' else Variant.DEVICE
746      return Variant.HOST if self._is_host_enabled else Variant.DEVICE
747
748    if not self._is_host_enabled and self.module_info.requires_device(info):
749      return DeviceTest(info, _select_variant(info), t_info.mainline_modules)
750
751    return DevicelessTest(info, _select_variant(info))
752
753  def _guess_test_type_for_missing_module(
754      self, t_info: test_info.TestInfo
755  ) -> Test:
756    """Determine the test type (device or deviceless) without module info."""
757    if (
758        not self._is_host_enabled
759        and t_info.get_supported_exec_mode() != constants.DEVICELESS_TEST
760    ):
761      return DeviceTest(None, Variant.DEVICE, t_info.mainline_modules)
762
763    return DevicelessTest(None, Variant.HOST)
764
765  def _get_host_framework_targets(self) -> Set[str]:
766    """Get the build targets for all the existing jars under host framework.
767
768    Returns:
769        A set of build target name under $(ANDROID_HOST_OUT)/framework.
770    """
771    host_targets = set()
772    if not self.module_info:
773      return host_targets
774
775    framework_host_dir = Path(
776        os.environ.get(constants.ANDROID_HOST_OUT)
777    ).joinpath('framework')
778    if framework_host_dir.is_dir():
779      jars = framework_host_dir.glob('*.jar')
780      for jar in jars:
781        if self.module_info.is_module(jar.stem):
782          host_targets.add(jar.stem)
783      logging.debug('Found exist host framework target:%s', host_targets)
784    return host_targets
785
786  def _parse_extra_args(self, test_infos, extra_args):
787    """Convert the extra args into something tf can understand.
788
789    Args:
790        extra_args: Dict of args
791
792    Returns:
793        Tuple of args to append and args not supported.
794    """
795    args_to_append, args_not_supported = extra_args_to_tf_args(
796        extra_args, self.module_info
797    )
798
799    # Set exclude instant app annotation for non-instant mode run.
800    if constants.INSTANT not in extra_args and self._has_instant_app_config(
801        test_infos, self.module_info
802    ):
803      args_to_append.append(constants.TF_TEST_ARG)
804      args_to_append.append(
805          '{tf_class}:{option_name}:{option_value}'.format(
806              tf_class=constants.TF_AND_JUNIT_CLASS,
807              option_name=constants.TF_EXCLUDE_ANNOTATE,
808              option_value=constants.INSTANT_MODE_ANNOTATE,
809          )
810      )
811    # Force append --enable-parameterized-modules if args_to_append has
812    # --module-parameter in args_to_append
813    if constants.TF_MODULE_PARAMETER in args_to_append:
814      if constants.TF_ENABLE_PARAMETERIZED_MODULES not in args_to_append:
815        args_to_append.append(constants.TF_ENABLE_PARAMETERIZED_MODULES)
816    # If all the test config has config with auto enable parameter, force
817    # exclude those default parameters(ex: instant_app, secondary_user)
818    # TODO: (b/228433541) Remove the limitation after the root cause fixed.
819    if len(test_infos) <= 1 and self._is_all_tests_parameter_auto_enabled(
820        test_infos
821    ):
822      if constants.TF_ENABLE_PARAMETERIZED_MODULES not in args_to_append:
823        args_to_append.append(constants.TF_ENABLE_PARAMETERIZED_MODULES)
824        for exclude_parameter in constants.DEFAULT_EXCLUDE_PARAS:
825          args_to_append.append('--exclude-module-parameters')
826          args_to_append.append(exclude_parameter)
827    return args_to_append, args_not_supported
828
829  def generate_run_commands(self, test_infos, extra_args, port=None):
830    """Generate a single run command from TestInfos.
831
832    Args:
833        test_infos: A list of TestInfo instances.
834        extra_args: A Dict of extra args to append.
835        port: Optional. An int of the port number to send events to. If None,
836          then subprocess reporter in TF won't try to connect.
837
838    Returns:
839        A list that contains the string of atest tradefed run command.
840        Only one command is returned.
841    """
842    if extra_args.get(constants.USE_TF_MIN_BASE_TEMPLATE):
843      self.run_cmd_dict['template'] = self._TF_LOCAL_MIN
844    else:
845      self.run_cmd_dict['template'] = (
846          self._TF_DEVICELESS_TEST_TEMPLATE
847          if extra_args.get(constants.HOST)
848          else self._TF_DEVICE_TEST_TEMPLATE
849      )
850
851    args = self._create_test_args(test_infos)
852
853    # Create a copy of args as more args could be added to the list.
854    test_args = list(args)
855    if port:
856      test_args.extend(['--subprocess-report-port', str(port)])
857    if extra_args.get(constants.INVOCATION_ID, None):
858      test_args.append(
859          '--invocation-data invocation_id=%s'
860          % extra_args[constants.INVOCATION_ID]
861      )
862    if extra_args.get(constants.WORKUNIT_ID, None):
863      test_args.append(
864          '--invocation-data work_unit_id=%s'
865          % extra_args[constants.WORKUNIT_ID]
866      )
867    if extra_args.get(constants.LOCAL_BUILD_ID, None):
868      # TODO: (b/207584685) Replace with TF local build solutions.
869      test_args.append('--use-stub-build true')
870      test_args.append(
871          '--stub-build-id %s' % extra_args[constants.LOCAL_BUILD_ID]
872      )
873      test_args.append(
874          '--stub-build-target %s' % extra_args[constants.BUILD_TARGET]
875      )
876    for info in test_infos:
877      if atest_utils.get_test_and_mainline_modules(info.test_name):
878        # TODO(b/253641058) Remove this once mainline module
879        # binaries are stored under testcase directory.
880        if not extra_args.get(constants.DRY_RUN):
881          self._copy_mainline_module_binary(info.mainline_modules)
882        test_args.append(constants.TF_ENABLE_MAINLINE_PARAMETERIZED_MODULES)
883        break
884    # For detailed logs, set TF options log-level/log-level-display as
885    # 'VERBOSE' by default.
886    log_level = 'VERBOSE'
887    test_args.extend(['--log-level-display', log_level])
888    test_args.extend(['--log-level', log_level])
889
890    # Set no-early-device-release by default to speed up TF teardown time.
891    # TODO(b/300882567) remove this forever when it's the default behavor.
892    test_args.extend(['--no-early-device-release'])
893
894    args_to_add, args_not_supported = self._parse_extra_args(
895        test_infos, extra_args
896    )
897
898    # If multiple devices in test config, automatically append
899    # --replicate-parent-setup and --multi-device-count
900    device_count = atest_configs.GLOBAL_ARGS.device_count_config
901    if device_count and device_count > 1:
902      args_to_add.append('--replicate-parent-setup')
903      args_to_add.append('--multi-device-count')
904      args_to_add.append(str(device_count))
905      os.environ.pop(constants.ANDROID_SERIAL, None)
906    else:
907      # TODO(b/122889707) Remove this after finding the root cause.
908      env_serial = os.environ.get(constants.ANDROID_SERIAL)
909      # Use the env variable ANDROID_SERIAL if it's set by user but only
910      # when the target tests are not deviceless tests.
911      if (
912          env_serial
913          and '--serial' not in args_to_add
914          and '-n' not in args_to_add
915      ):
916        args_to_add.append('--serial')
917        args_to_add.append(env_serial)
918
919    test_args.extend(args_to_add)
920    if args_not_supported:
921      atest_utils.print_and_log_info(
922          '%s does not support the following args %s',
923          self.EXECUTABLE,
924          args_not_supported,
925      )
926
927    # Only need to check one TestInfo to determine if the tests are
928    # configured in TEST_MAPPING.
929    for_test_mapping = test_infos and test_infos[0].from_test_mapping
930    if is_log_upload_enabled(extra_args):
931      test_args.extend(atest_utils.get_result_server_args(for_test_mapping))
932    self.run_cmd_dict['args'] = ' '.join(test_args)
933    self.run_cmd_dict['tf_customize_template'] = (
934        self._extract_customize_tf_templates(extra_args, test_infos)
935    )
936
937    # By default using ATestFileSystemLogSaver no matter what running under
938    # aosp or internal branches. Only switch using google log saver if user
939    # tend to upload test result to AnTS which could be detected by the
940    # invocation_id in extra args.
941    if is_log_upload_enabled(extra_args):
942      self.use_google_log_saver()
943
944    run_commands = [self._RUN_CMD.format(**self.run_cmd_dict)]
945    logging.debug('TF test runner generated run commands %s', run_commands)
946    return run_commands
947
948  def _flatten_test_infos(self, test_infos):
949    """Sort and group test_infos by module_name and sort and group filters
950
951    by class name.
952
953        Example of three test_infos in a set:
954            Module1, {(classA, {})}
955            Module1, {(classB, {Method1})}
956            Module1, {(classB, {Method2}}
957        Becomes a set with one element:
958            Module1, {(ClassA, {}), (ClassB, {Method1, Method2})}
959        Where:
960              Each line is a test_info namedtuple
961              {} = Frozenset
962              () = TestFilter namedtuple
963
964    Args:
965        test_infos: A list of TestInfo namedtuples.
966
967    Returns:
968        A list of TestInfos flattened.
969    """
970    results = []
971    for module, group in atest_utils.sort_and_group(
972        test_infos, lambda x: x.test_name
973    ):
974
975      # module is a string, group is a generator of grouped TestInfos.
976      # Module Test, so flatten test_infos:
977      no_filters = False
978      filters = set()
979      test_runner = None
980      test_finder = None
981      build_targets = set()
982      data = {}
983      module_args = []
984      for test_info_i in group:
985        data.update(test_info_i.data)
986        # Extend data with constants.TI_MODULE_ARG instead of
987        # overwriting.
988        module_args.extend(test_info_i.data.get(constants.TI_MODULE_ARG, []))
989        test_runner = test_info_i.test_runner
990        test_finder = test_info_i.test_finder
991        build_targets |= test_info_i.build_targets
992        test_filters = test_info_i.data.get(constants.TI_FILTER)
993        if not test_filters or no_filters:
994          # test_info wants whole module run, so hardcode no filters.
995          no_filters = True
996          filters = set()
997          continue
998        filters |= test_filters
999      if module_args:
1000        data[constants.TI_MODULE_ARG] = module_args
1001      data[constants.TI_FILTER] = self.flatten_test_filters(filters)
1002      results.append(
1003          test_info.TestInfo(
1004              test_name=module,
1005              test_runner=test_runner,
1006              test_finder=test_finder,
1007              build_targets=build_targets,
1008              data=data,
1009          )
1010      )
1011    return results
1012
1013  @staticmethod
1014  def flatten_test_filters(filters):
1015    """Sort and group test_filters by class_name.
1016
1017        Example of three test_filters in a frozenset:
1018            classA, {}
1019            classB, {Method1}
1020            classB, {Method2}
1021        Becomes a frozenset with these elements:
1022            classA, {}
1023            classB, {Method1, Method2}
1024        Where:
1025            Each line is a TestFilter namedtuple
1026            {} = Frozenset
1027
1028    Args:
1029        filters: A frozenset of test_filters.
1030
1031    Returns:
1032        A frozenset of test_filters flattened.
1033    """
1034    results = set()
1035    for class_name, group in atest_utils.sort_and_group(
1036        filters, lambda x: x.class_name
1037    ):
1038
1039      # class_name is a string, group is a generator of TestFilters
1040      assert class_name is not None
1041      methods = set()
1042      for test_filter in group:
1043        if not test_filter.methods:
1044          # Whole class should be run
1045          methods = set()
1046          break
1047        methods |= test_filter.methods
1048      results.add(test_info.TestFilter(class_name, frozenset(methods)))
1049    return frozenset(results)
1050
1051  def _is_all_tests_parameter_auto_enabled(self, test_infos):
1052    """Check if all the test infos are parameter auto enabled.
1053
1054    Args:
1055        test_infos: A set of TestInfo instances.
1056
1057    Returns: True if all tests are parameter auto enabled, False otherwise.
1058    """
1059    for info in test_infos:
1060      if not self._is_parameter_auto_enabled_cfg(info, self.module_info):
1061        return False
1062    return True
1063
1064  def _create_test_args(self, test_infos):
1065    """Compile TF command line args based on the given test infos.
1066
1067    Args:
1068        test_infos: A list of TestInfo instances.
1069
1070    Returns: A list of TF arguments to run the tests.
1071    """
1072    args = []
1073    if not test_infos:
1074      return []
1075
1076    if atest_configs.GLOBAL_ARGS.group_test:
1077      test_infos = self._flatten_test_infos(test_infos)
1078
1079    has_integration_test = False
1080
1081    # Because current --include-filter arg will not working if ATest pass
1082    # both --module and --include-filter to TF, only test by --module will
1083    # be run. Make a check first, only use --module if all tests are all
1084    # parameter auto enabled.
1085    # Only auto-enable the parameter if there's only one test.
1086    # TODO: (b/228433541) Remove the limitation after the root cause fixed.
1087    use_module_arg = False
1088    if len(test_infos) <= 1:
1089      use_module_arg = self._is_all_tests_parameter_auto_enabled(test_infos)
1090
1091    for info in test_infos:
1092      # Integration test exists in TF's jar, so it must have the option
1093      # if it's integration finder.
1094      if info.test_finder in _INTEGRATION_FINDERS:
1095        has_integration_test = True
1096      # For non-parameterize test module, use --include-filter, but for
1097      # tests which have auto enable parameterize config use --module
1098      # instead.
1099      if use_module_arg and self._is_parameter_auto_enabled_cfg(
1100          info, self.module_info
1101      ):
1102        args.extend([constants.TF_MODULE_FILTER, info.test_name])
1103      else:
1104        args.extend([constants.TF_INCLUDE_FILTER, info.test_name])
1105      for option in info.data.get(constants.TI_MODULE_ARG, []):
1106        if constants.TF_INCLUDE_FILTER_OPTION == option[0]:
1107          suite_filter = constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
1108              test_name=info.test_name, option_value=option[1]
1109          )
1110          args.extend([constants.TF_INCLUDE_FILTER, suite_filter])
1111        elif constants.TF_EXCLUDE_FILTER_OPTION == option[0]:
1112          suite_filter = constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
1113              test_name=info.test_name, option_value=option[1]
1114          )
1115          args.extend([constants.TF_EXCLUDE_FILTER, suite_filter])
1116        else:
1117          module_arg = constants.TF_MODULE_ARG_VALUE_FMT.format(
1118              test_name=info.test_name,
1119              option_name=option[0],
1120              option_value=option[1],
1121          )
1122          args.extend([constants.TF_MODULE_ARG, module_arg])
1123
1124    # Add ATest include filter
1125    args.extend(get_include_filter(test_infos))
1126
1127    # TODO (b/141090547) Pass the config path to TF to load configs.
1128    # Compile option in TF if finder is not INTEGRATION or not set.
1129    if not has_integration_test:
1130      args.append(constants.TF_SKIP_LOADING_CONFIG_JAR)
1131    return args
1132
1133  def _extract_rerun_options(self, extra_args):
1134    """Extract rerun options to a string for output.
1135
1136    Args:
1137        extra_args: Dict of extra args for test runners to use.
1138
1139    Returns: A string of rerun options.
1140    """
1141    extracted_options = [
1142        '{} {}'.format(arg, extra_args[arg])
1143        for arg in extra_args
1144        if arg in self._RERUN_OPTION_GROUP
1145    ]
1146    return ' '.join(extracted_options)
1147
1148  def _extract_customize_tf_templates(self, extra_args, test_infos):
1149    """Extract tradefed template options to a string for output.
1150
1151    Args:
1152        extra_args: Dict of extra args for test runners to use.
1153        test_infos: A set of TestInfo instances.
1154
1155    Returns: A string of tradefed template options.
1156    """
1157    tf_templates = extra_args.get(constants.TF_TEMPLATE, [])
1158    tf_template_keys = [i.split('=')[0] for i in tf_templates]
1159    for info in test_infos:
1160      if (
1161          info.aggregate_metrics_result
1162          and 'metric_post_processor' not in tf_template_keys
1163      ):
1164        template_key = 'metric_post_processor'
1165        template_value = (
1166            'google/template/postprocessors/metric-file-aggregate-disabled'
1167        )
1168        tf_templates.append(f'{template_key}={template_value}')
1169    return ' '.join(['--template:map %s' % x for x in tf_templates])
1170
1171  def _handle_log_associations(self, event_handlers):
1172    """Handle TF's log associations information data.
1173
1174    log_association dict:
1175    {'loggedFile': '/tmp/serial-util11375755456514097276.ser',
1176     'dataName': 'device_logcat_setup_127.0.0.1:58331',
1177     'time': 1602038599.856113},
1178
1179    Args:
1180        event_handlers: Dict of {socket_object:EventHandler}.
1181    """
1182    log_associations = []
1183    for _, event_handler in event_handlers.items():
1184      if event_handler.log_associations:
1185        log_associations += event_handler.log_associations
1186    device_test_end_log_time = ''
1187    device_teardown_log_time = ''
1188    for log_association in log_associations:
1189      if 'device_logcat_test' in log_association.get('dataName', ''):
1190        device_test_end_log_time = log_association.get('time')
1191      if 'device_logcat_teardown' in log_association.get('dataName', ''):
1192        device_teardown_log_time = log_association.get('time')
1193    if device_test_end_log_time and device_teardown_log_time:
1194      teardowntime = float(device_teardown_log_time) - float(
1195          device_test_end_log_time
1196      )
1197      logging.debug('TF logcat teardown time=%s seconds.', teardowntime)
1198      metrics.LocalDetectEvent(
1199          detect_type=DetectType.TF_TEARDOWN_LOGCAT, result=int(teardowntime)
1200      )
1201
1202  @staticmethod
1203  def _has_instant_app_config(test_infos, mod_info):
1204    """Check if one of the input tests defined instant app mode in config.
1205
1206    Args:
1207        test_infos: A set of TestInfo instances.
1208        mod_info: ModuleInfo object.
1209
1210    Returns: True if one of the tests set up instant app mode.
1211    """
1212    for tinfo in test_infos:
1213      test_config, _ = test_finder_utils.get_test_config_and_srcs(
1214          tinfo, mod_info
1215      )
1216      if test_config:
1217        parameters = atest_utils.get_config_parameter(test_config)
1218        if constants.TF_PARA_INSTANT_APP in parameters:
1219          return True
1220    return False
1221
1222  @staticmethod
1223  def _is_parameter_auto_enabled_cfg(tinfo, mod_info):
1224    """Check if input tests contains auto enable support parameters.
1225
1226    Args:
1227        test_infos: A set of TestInfo instances.
1228        mod_info: ModuleInfo object.
1229
1230    Returns: True if input test has parameter setting which is not in the
1231             exclude list.
1232    """
1233    test_config, _ = test_finder_utils.get_test_config_and_srcs(tinfo, mod_info)
1234    if test_config:
1235      parameters = atest_utils.get_config_parameter(test_config)
1236      if (
1237          parameters
1238          - constants.DEFAULT_EXCLUDE_PARAS
1239          - constants.DEFAULT_EXCLUDE_NOT_PARAS
1240      ):
1241        return True
1242    return False
1243
1244  def _handle_native_tests(self, test_infos):
1245    """Handling some extra tasks for running native tests from tradefed.
1246
1247    Args:
1248        test_infos: A set of TestInfo instances.
1249    """
1250    for tinfo in test_infos:
1251      test_config, _ = test_finder_utils.get_test_config_and_srcs(
1252          tinfo, self.module_info
1253      )
1254      if test_config:
1255        module_name, device_path = atest_utils.get_config_gtest_args(
1256            test_config
1257        )
1258        if module_name and device_path:
1259          atest_utils.copy_native_symbols(module_name, device_path)
1260
1261  # TODO(b/253641058) remove copying files once mainline module
1262  # binaries are stored under testcase directory.
1263  def _copy_mainline_module_binary(self, mainline_modules):
1264    """Copies mainline module binaries to out/dist/mainline_modules_{arch}
1265
1266    Copies the mainline module binaries to the location that
1267    MainlineModuleHandler in TF expects since there is no way to
1268    explicitly tweak the search path.
1269
1270    Args:
1271        mainline_modules: A list of mainline modules.
1272    """
1273    config = atest_utils.get_android_config()
1274    arch = config.get('TARGET_ARCH')
1275    dest_dir = atest_utils.DIST_OUT_DIR.joinpath(f'mainline_modules_{arch}')
1276    dest_dir.mkdir(parents=True, exist_ok=True)
1277
1278    for module in mainline_modules:
1279      target_module_info = self.module_info.get_module_info(module)
1280      installed_paths = target_module_info[constants.MODULE_INSTALLED]
1281
1282      for installed_path in installed_paths:
1283        file_name = Path(installed_path).name
1284        dest_path = Path(dest_dir).joinpath(file_name)
1285        if dest_path.exists():
1286          atest_utils.colorful_print(
1287              'Replacing APEX in %s with %s' % (dest_path, installed_path),
1288              constants.CYAN,
1289          )
1290          logging.debug(
1291              'deleting the old file: %s and copy a new binary', dest_path
1292          )
1293          dest_path.unlink()
1294        shutil.copyfile(installed_path, dest_path)
1295
1296        break
1297
1298  def use_google_log_saver(self):
1299    """Replace the original log saver to google log saver."""
1300    self.log_args.update({
1301        'log_root_option_name': constants.GOOGLE_LOG_SAVER_LOG_ROOT_OPTION_NAME,
1302        'log_ext_option': constants.GOOGLE_LOG_SAVER_EXT_OPTION,
1303    })
1304    self.run_cmd_dict.update({
1305        'log_saver': constants.GOOGLE_LOG_SAVER,
1306        'log_args': self._LOG_ARGS.format(**self.log_args),
1307    })
1308
1309
1310def is_log_upload_enabled(extra_args: Dict[str, Any]) -> bool:
1311  """Check if input extra_args include google log saver related args.
1312
1313  Args:
1314      extra_args: Dict of args.
1315  """
1316  return bool(extra_args.get(constants.INVOCATION_ID, None))
1317
1318
1319def generate_annotation_filter_args(
1320    arg_value: Any,
1321    mod_info: module_info.ModuleInfo,
1322    test_infos: List[test_info.TestInfo],
1323) -> List[str]:
1324  """Generate TF annotation filter arguments.
1325
1326  Args:
1327      arg_value: Argument value for annotation filter.
1328      mod_info: ModuleInfo object.
1329      test_infos: A set of TestInfo instances.
1330
1331  Returns:
1332      List of TF annotation filter arguments.
1333  """
1334  annotation_filter_args = []
1335  for info in test_infos:
1336    test_name = info.test_name
1337    for keyword in arg_value:
1338      annotation = atest_utils.get_full_annotation_class_name(
1339          mod_info.get_module_info(test_name), keyword
1340      )
1341      if annotation:
1342        module_arg = constants.TF_MODULE_ARG_VALUE_FMT.format(
1343            test_name=test_name,
1344            option_name=constants.INCLUDE_ANNOTATION,
1345            option_value=annotation,
1346        )
1347        annotation_filter_args.extend([constants.TF_MODULE_ARG, module_arg])
1348      atest_utils.print_and_log_error(
1349          atest_utils.mark_red(f'Cannot find similar annotation: {keyword}')
1350      )
1351  return annotation_filter_args
1352
1353
1354def extra_args_to_tf_args(
1355    extra_args: Dict[str, Any], mod_info: module_info.ModuleInfo = None
1356) -> Tuple[Dict[str, Any], Dict[str, Any]]:
1357  """Convert the extra args into atest_tf_test_runner supported args.
1358
1359  Args:
1360      extra_args: Dict of args
1361      mod_info: ModuleInfo object.
1362
1363  Returns:
1364      Tuple of ARGS that atest_tf supported and not supported.
1365  """
1366  supported_args = []
1367  unsupported_args = []
1368
1369  def constant_list(*value):
1370    return lambda *_: value
1371
1372  # pylint: disable=unused-argument
1373  def print_message(message):
1374    def inner(*args):
1375      print(message)
1376      return []
1377
1378    return inner
1379
1380  # Mapping supported TF arguments to the processing function.
1381  supported_tf_args = dict({
1382      constants.WAIT_FOR_DEBUGGER: constant_list('--wait-for-debugger'),
1383      constants.DISABLE_INSTALL: constant_list('--disable-target-preparers'),
1384      constants.SERIAL: lambda arg_value: [
1385          j for d in arg_value for j in ('--serial', d)
1386      ],
1387      constants.SHARDING: lambda arg_value: ['--shard-count', str(arg_value)],
1388      constants.DISABLE_TEARDOWN: constant_list('--disable-teardown'),
1389      constants.HOST: constant_list(
1390          '-n', '--prioritize-host-config', '--skip-host-arch-check'
1391      ),
1392      constants.CUSTOM_ARGS:
1393      # custom args value is a list.
1394      lambda arg_value: arg_value,
1395      constants.ALL_ABI: constant_list('--all-abi'),
1396      constants.INSTANT: constant_list(
1397          constants.TF_ENABLE_PARAMETERIZED_MODULES,
1398          constants.TF_MODULE_PARAMETER,
1399          'instant_app',
1400      ),
1401      constants.USER_TYPE: lambda arg_value: [
1402          constants.TF_ENABLE_PARAMETERIZED_MODULES,
1403          '--enable-optional-parameterization',
1404          constants.TF_MODULE_PARAMETER,
1405          str(arg_value),
1406      ],
1407      constants.ITERATIONS: lambda arg_value: [
1408          '--retry-strategy',
1409          constants.ITERATIONS,
1410          '--max-testcase-run-count',
1411          str(arg_value),
1412      ],
1413      constants.RERUN_UNTIL_FAILURE: lambda arg_value: [
1414          '--retry-strategy',
1415          constants.RERUN_UNTIL_FAILURE,
1416          '--max-testcase-run-count',
1417          str(arg_value),
1418      ],
1419      constants.RETRY_ANY_FAILURE: lambda arg_value: [
1420          '--retry-strategy',
1421          constants.RETRY_ANY_FAILURE,
1422          '--max-testcase-run-count',
1423          str(arg_value),
1424      ],
1425      constants.COLLECT_TESTS_ONLY: constant_list('--collect-tests-only'),
1426      constants.TF_DEBUG: print_message('Please attach process to your IDE...'),
1427      constants.ANNOTATION_FILTER: generate_annotation_filter_args,
1428      constants.TEST_FILTER: lambda arg_value: [
1429          '--test-arg',
1430          (
1431              'com.android.tradefed.testtype.AndroidJUnitTest:'
1432              f'include-filter:{arg_value}'
1433          ),
1434          '--test-arg',
1435          (
1436              'com.android.tradefed.testtype.GTest:native-test-flag:'
1437              f'--gtest_filter={arg_value}'
1438          ),
1439          '--test-arg',
1440          (
1441              'com.android.tradefed.testtype.HostGTest:native-test-flag:'
1442              f'--gtest_filter={arg_value}'
1443          ),
1444      ],
1445      constants.TEST_TIMEOUT: lambda arg_value: [
1446          '--test-arg',
1447          (
1448              'com.android.tradefed.testtype.AndroidJUnitTest:'
1449              f'shell-timeout:{arg_value}'
1450          ),
1451          '--test-arg',
1452          (
1453              'com.android.tradefed.testtype.AndroidJUnitTest:'
1454              f'test-timeout:{arg_value}'
1455          ),
1456          '--test-arg',
1457          (
1458              'com.android.tradefed.testtype.HostGTest:'
1459              f'native-test-timeout:{arg_value}'
1460          ),
1461          '--test-arg',
1462          (
1463              'com.android.tradefed.testtype.GTest:'
1464              f'native-test-timeout:{arg_value}'
1465          ),
1466          '--test-arg',
1467          (
1468              'com.android.compatibility.testtype.LibcoreTest:'
1469              f'test-timeout:{arg_value}'
1470          ),
1471      ],
1472      constants.COVERAGE: lambda _: coverage.tf_args(mod_info),
1473  })
1474
1475  for arg in extra_args:
1476    if arg in supported_tf_args:
1477      tf_args = supported_tf_args[arg](extra_args[arg])
1478      if tf_args:
1479        supported_args.extend(tf_args)
1480      continue
1481
1482    if arg in (
1483        constants.TF_TEMPLATE,
1484        constants.INVOCATION_ID,
1485        constants.WORKUNIT_ID,
1486        constants.REQUEST_UPLOAD_RESULT,
1487        constants.DISABLE_UPLOAD_RESULT,
1488        constants.LOCAL_BUILD_ID,
1489        constants.BUILD_TARGET,
1490        constants.DRY_RUN,
1491        constants.DEVICE_ONLY,
1492    ):
1493      continue
1494    unsupported_args.append(arg)
1495  return supported_args, unsupported_args
1496
1497
1498def get_include_filter(test_infos: List[test_info.TestInfo]) -> List[str]:
1499  """Generate a list of tradefed filter argument from TestInfos.
1500
1501  Args:
1502      test_infos: a List of TestInfo object.
1503
1504  The include filter pattern looks like:
1505      --atest-include-filter <module-name>:<include-filter-value>
1506
1507  Returns:
1508      List of Tradefed command args.
1509  """
1510  instrumentation_filters = []
1511  tf_args = []
1512  for info in test_infos:
1513    filters = []
1514    for test_info_filter in info.data.get(constants.TI_FILTER, []):
1515      filters.extend(test_info_filter.to_list_of_tf_strings())
1516
1517    for test_filter in filters:
1518      filter_arg = constants.TF_ATEST_INCLUDE_FILTER_VALUE_FMT.format(
1519          test_name=info.test_name, test_filter=test_filter
1520      )
1521      tf_args.extend([constants.TF_ATEST_INCLUDE_FILTER, filter_arg])
1522
1523  return tf_args
1524
1525
1526@enum.unique
1527class Variant(enum.Enum):
1528  """The variant of a build module."""
1529
1530  NONE = ''
1531  HOST = 'host'
1532  DEVICE = 'target'
1533
1534  def __init__(self, suffix):
1535    self._suffix = suffix
1536
1537  @property
1538  def suffix(self) -> str:
1539    """The suffix without the 'dash' used to qualify build targets."""
1540    return self._suffix
1541
1542
1543@dataclasses.dataclass(frozen=True)
1544class Target:
1545  """A build target."""
1546
1547  module_name: str
1548  variant: Variant
1549
1550  def name(self) -> str:
1551    """The name to use on the command-line to build this target."""
1552    if not self.variant.suffix:
1553      return self.module_name
1554    return f'{self.module_name}-{self.variant.suffix}'
1555
1556
1557class Test(ABC):
1558  """A test that can be run."""
1559
1560  _DEFAULT_HARNESS_TARGETS = frozenset(
1561      [
1562          Target('atest-tradefed', Variant.HOST),
1563          Target('atest_script_help.sh', Variant.HOST),
1564          Target('atest_tradefed.sh', Variant.HOST),
1565          Target('tradefed', Variant.HOST),
1566      ]
1567      + [Target(t, Variant.HOST) for t in constants.GTF_TARGETS]
1568  )
1569
1570  def query_build_targets(self) -> Set[Target]:
1571    """Returns the list of build targets required to run this test."""
1572    build_targets = set()
1573    build_targets.update(self._get_harness_build_targets())
1574    build_targets.update(self._get_test_build_targets())
1575    return build_targets
1576
1577  @abstractmethod
1578  def query_runtime_targets(self) -> Set[Target]:
1579    """Returns the list of targets required during runtime."""
1580
1581  @abstractmethod
1582  def _get_test_build_targets(self) -> Set[Target]:
1583    """Returns the list of build targets of test and its dependencies."""
1584
1585  @abstractmethod
1586  def _get_harness_build_targets(self) -> Set[Target]:
1587    """Returns the list of build targets of test harness and its dependencies."""
1588
1589  @abstractmethod
1590  def requires_device(self) -> bool:
1591    """Returns true if the test requires a device, otherwise false."""
1592
1593  @abstractmethod
1594  def requires_device_update(self) -> bool:
1595    """Checks whether the test requires device update."""
1596
1597
1598class DeviceTest(Test):
1599  """A device test that can be run."""
1600
1601  def __init__(
1602      self, info: Dict[str, Any], variant: Variant, mainline_modules: Set[str]
1603  ):
1604
1605    self._info = info
1606    self._variant = variant
1607    self._mainline_modules = mainline_modules
1608
1609  def query_runtime_targets(self) -> Set[Target]:
1610    return self.query_build_targets() | _get_host_required_deps(self._info)
1611
1612  def requires_device(self):
1613    return True
1614
1615  def requires_device_update(self):
1616    # The test doesn't need device update as long as it's a unit test,
1617    # no matter if it's running on device or host.
1618    # Some tests (e.g. TF integration tests) do not have module info, and we
1619    # can't determine whether they require device update or not. So that we
1620    # treat them as they require device update to avoid disabling the device
1621    # update mistakenly.
1622    return not self._info or not module_info.ModuleInfo.is_unit_test(
1623        self._info)
1624
1625  def _get_test_build_targets(self) -> Set[Target]:
1626    module_name = self._info[constants.MODULE_INFO_ID]
1627    build_targets = set([Target(module_name, self._variant)])
1628    build_targets.update(_get_libs_deps(self._info, self._variant))
1629    build_targets.update(
1630        Target(m, Variant.NONE) for m in self._mainline_modules
1631    )
1632    return build_targets
1633
1634  def _get_harness_build_targets(self):
1635    build_targets = set(Test._DEFAULT_HARNESS_TARGETS)
1636    build_targets.update(
1637        set([
1638            Target('adb', Variant.HOST),
1639            Target('aapt', Variant.HOST),
1640            Target('aapt2', Variant.HOST),
1641            Target('compatibility-host-util', Variant.HOST),
1642        ])
1643    )
1644
1645    # Auto-generated Java tests use a module template that uses the Dalvik
1646    # test runner and requires the implementation jars. See
1647    # https://source.corp.google.com/android-internal/build/make/core/java_test_config_template.xml.
1648    # These dependencies should ideally be automatically added by the build
1649    # rule since Atest can fall out of sync otherwise.
1650    # TODO(b/284987354): Remove these targets once the build rule adds the
1651    # required deps.
1652    if _is_dalvik_test_module(self._info):
1653      build_targets.add(Target('cts-dalvik-host-test-runner', Variant.HOST))
1654      build_targets.add(Target('cts-dalvik-device-test-runner', Variant.DEVICE))
1655
1656    if 'vts' in self._info.get(constants.MODULE_COMPATIBILITY_SUITES, []):
1657      # Note that we do not include `compatibility-tradefed` which is
1658      # already included in the VTS harness.
1659      build_targets.add(Target('vts-core-tradefed-harness', Variant.HOST))
1660    else:
1661      build_targets.add(Target('compatibility-tradefed', Variant.HOST))
1662
1663    return build_targets
1664
1665
1666class DevicelessTest(Test):
1667
1668  def __init__(self, info: Dict[str, Any], variant: Variant):
1669    self._info = info
1670    self._variant = variant
1671
1672  def _get_test_build_targets(self) -> Set[Target]:
1673    module_name = self._info[constants.MODULE_INFO_ID]
1674    return set([Target(module_name, self._variant)])
1675
1676  def _get_harness_build_targets(self):
1677    build_targets = set(Test._DEFAULT_HARNESS_TARGETS)
1678    build_targets.update(
1679        set([
1680            # TODO(b/277116853): Remove the adb dependency for deviceless tests.
1681            Target('adb', Variant.HOST),
1682        ])
1683    )
1684    return build_targets
1685
1686  def query_runtime_targets(self) -> Set[Target]:
1687    return self.query_build_targets()
1688
1689  def requires_device(self):
1690    return False
1691
1692  def requires_device_update(self):
1693    return False
1694
1695
1696def _get_libs_deps(info: Dict[str, Any], variant: Variant) -> Set[Target]:
1697
1698  # We only need the runtime dependencies with host variant since TradeFed
1699  # won't push any runtime dependencies to the test device and the runtime
1700  # dependencies with device variant should already exist on the test device.
1701  if variant != Variant.HOST:
1702    return set()
1703
1704  deps = set()
1705  deps.update([Target(m, variant) for m in info.get(constants.MODULE_LIBS, [])])
1706
1707  return deps
1708
1709
1710def _get_host_required_deps(info: Dict[str, Any]) -> Set[Target]:
1711
1712  deps = set()
1713  deps.update(
1714      Target(m, Variant.HOST) for m in info.get(constants.MODULE_HOST_DEPS, [])
1715  )
1716
1717  return deps
1718
1719
1720def _is_dalvik_test_module(info: Dict[str, Any]) -> bool:
1721  return 'JAVA_LIBRARIES' in info.get(
1722      constants.MODULE_CLASS, []
1723  ) and True in info.get(constants.MODULE_AUTO_TEST_CONFIG, [])
1724