• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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=line-too-long
18# pylint: disable=too-many-lines
19
20from __future__ import print_function
21
22import json
23import logging
24import os
25import re
26import select
27import shutil
28import socket
29
30from functools import partial
31from pathlib import Path
32from typing import Any, Dict, List, Set, Tuple
33
34from atest import atest_configs
35from atest import atest_error
36from atest import atest_utils
37from atest import constants
38from atest import module_info
39from atest import result_reporter
40
41from atest.atest_enum import DetectType, ExitCode
42from atest.coverage import coverage
43from atest.logstorage import atest_gcp_utils
44from atest.logstorage import logstorage_utils
45from atest.metrics import metrics
46from atest.test_finders import test_finder_utils
47from atest.test_finders import test_info
48from atest.test_runners import test_runner_base as trb
49from atest.test_runners.event_handler import EventHandler
50
51POLL_FREQ_SECS = 10
52SOCKET_HOST = '127.0.0.1'
53SOCKET_QUEUE_MAX = 1
54SOCKET_BUFFER = 4096
55SELECT_TIMEOUT = 0.5
56
57# Socket Events of form FIRST_EVENT {JSON_DATA}\nSECOND_EVENT {JSON_DATA}
58# EVENT_RE has groups for the name and the data. "." does not match \n.
59EVENT_RE = re.compile(r'\n*(?P<event_name>[A-Z_]+) (?P<json_data>{.*})(?=\n|.)*')
60
61# Remove aapt from build dependency, use prebuilt version instead.
62EXEC_DEPENDENCIES = ('adb', 'fastboot')
63
64LOG_FOLDER_NAME = 'log'
65
66_INTEGRATION_FINDERS = frozenset(['', 'INTEGRATION', 'INTEGRATION_FILE_PATH'])
67
68# AAPT binary name
69_AAPT = 'aapt'
70
71# The exist code mapping of tradefed.
72_TF_EXIT_CODE = [
73    'NO_ERROR',
74    'CONFIG_EXCEPTION',
75    'NO_BUILD',
76    'DEVICE_UNRESPONSIVE',
77    'DEVICE_UNAVAILABLE',
78    'FATAL_HOST_ERROR',
79    'THROWABLE_EXCEPTION',
80    'NO_DEVICE_ALLOCATED',
81    'WRONG_JAVA_VERSION']
82
83MAINLINE_LOCAL_DOC = 'go/mainline-local-build'
84
85class TradeFedExitError(Exception):
86    """Raised when TradeFed exists before test run has finished."""
87    def __init__(self, exit_code):
88        super().__init__()
89        self.exit_code = exit_code
90
91    def __str__(self):
92        tf_error_reason = self._get_exit_reason(self.exit_code)
93        return (f'TradeFed subprocess exited early with exit code='
94                f'{self.exit_code}({tf_error_reason}).')
95
96    def _get_exit_reason(self, exit_code):
97        if 0 < exit_code < len(_TF_EXIT_CODE):
98            return atest_utils.colorize(_TF_EXIT_CODE[exit_code], constants.RED)
99        return 'Unknown exit status'
100
101class AtestTradefedTestRunner(trb.TestRunnerBase):
102    """TradeFed Test Runner class."""
103    NAME = 'AtestTradefedTestRunner'
104    EXECUTABLE = 'atest_tradefed.sh'
105    _TF_TEMPLATE = 'template/atest_local_min'
106    # Use --no-enable-granular-attempts to control reporter replay behavior.
107    # TODO(b/142630648): Enable option enable-granular-attempts
108    # in sharding mode.
109    _LOG_ARGS = ('--logcat-on-failure --{log_root_option_name}={log_path} '
110                 '{log_ext_option} '
111                 '--no-enable-granular-attempts '
112                 '--proto-output-file={proto_path}')
113    _RUN_CMD = ('{env} {exe} {template} '
114                '--template:map test=atest '
115                '--template:map log_saver={log_saver} '
116                '{tf_customize_template} {log_args} {args}')
117    _BUILD_REQ = {'tradefed-core'}
118    _RERUN_OPTION_GROUP = [constants.ITERATIONS,
119                           constants.RERUN_UNTIL_FAILURE,
120                           constants.RETRY_ANY_FAILURE]
121
122    def __init__(self, results_dir: str,
123                 mod_info: module_info.ModuleInfo=None, **kwargs):
124        """Init stuff for base class."""
125        super().__init__(results_dir, **kwargs)
126        self.module_info = mod_info
127        self.log_path = os.path.join(results_dir, LOG_FOLDER_NAME)
128        # (b/275537997) results_dir could be '' in test_runner_handler; only
129        # mkdir when it is invoked by run_tests.
130        if results_dir:
131            Path(self.log_path).mkdir(parents=True, exist_ok=True)
132        log_args = {'log_root_option_name': constants.LOG_ROOT_OPTION_NAME,
133                    'log_ext_option': constants.LOG_SAVER_EXT_OPTION,
134                    'log_path': self.log_path,
135                    'proto_path': os.path.join(
136                        self.results_dir,
137                        constants.ATEST_TEST_RECORD_PROTO)}
138        self.run_cmd_dict = {'env': '',
139                             'exe': self.EXECUTABLE,
140                             'template': self._TF_TEMPLATE,
141                             'log_saver': constants.ATEST_TF_LOG_SAVER,
142                             'tf_customize_template': '',
143                             'args': '',
144                             'log_args': self._LOG_ARGS.format(**log_args)}
145        if kwargs.get('extra_args', {}).get(constants.LD_LIBRARY_PATH, False):
146            self.run_cmd_dict.update({'env': self._get_ld_library_path()})
147        self.is_verbose = logging.getLogger().isEnabledFor(logging.DEBUG)
148        self.root_dir = os.environ.get(constants.ANDROID_BUILD_TOP)
149
150    def _get_ld_library_path(self) -> str:
151        """Get the corresponding LD_LIBRARY_PATH string for running TF.
152
153        This method will insert $ANDROID_HOST_OUT/{lib,lib64} to LD_LIBRARY_PATH
154        and returns the updated LD_LIBRARY_PATH.
155
156        Returns:
157            Strings for the environment passed to TF. Currently only
158            LD_LIBRARY_PATH for TF to load the correct local shared libraries.
159        """
160        out_dir = os.environ.get(constants.ANDROID_HOST_OUT, '')
161        # From b/188179058, if a 64bit tests, it will break the tests due to the
162        # elf format is not 64bit for the lib path. But for b/160741384, it is
163        # ok to load lib path first. Change the lib_dirs sequence to lib64 first
164        # due to ATest by default only testing the main abi and even a 32bit
165        # only target the lib64 folder is actually not exist.
166        lib_dirs = ['lib64', 'lib']
167        path = ':'.join([os.path.join(out_dir, dir) for dir in lib_dirs])
168        return f'LD_LIBRARY_PATH={path}:{os.getenv("LD_LIBRARY_PATH", "")}'
169
170    def _try_set_gts_authentication_key(self):
171        """Set GTS authentication key if it is available or exists.
172
173        Strategy:
174            Get APE_API_KEY from os.environ:
175                - If APE_API_KEY is already set by user -> do nothing.
176            Get the APE_API_KEY from constants:
177                - If the key file exists -> set to env var.
178            If APE_API_KEY isn't set and the key file doesn't exist:
179                - Warn user some GTS tests may fail without authentication.
180        """
181        if os.environ.get('APE_API_KEY'):
182            logging.debug('APE_API_KEY is set by developer.')
183            return
184        ape_api_key = constants.GTS_GOOGLE_SERVICE_ACCOUNT
185        key_path = os.path.join(self.root_dir, ape_api_key)
186        if ape_api_key and os.path.exists(key_path):
187            logging.debug('Set APE_API_KEY: %s', ape_api_key)
188            os.environ['APE_API_KEY'] = key_path
189        else:
190            logging.debug('APE_API_KEY not set, some GTS tests may fail'
191                          ' without authentication.')
192
193    def run_tests(self, test_infos, extra_args, reporter):
194        """Run the list of test_infos. See base class for more.
195
196        Args:
197            test_infos: A list of TestInfos.
198            extra_args: Dict of extra args to add to test run.
199            reporter: An instance of result_report.ResultReporter.
200
201        Returns:
202            0 if tests succeed, non-zero otherwise.
203        """
204        reporter.log_path = self.log_path
205        reporter.rerun_options = self._extract_rerun_options(extra_args)
206        # Set google service key if it's available or found before
207        # running tests.
208        self._try_set_gts_authentication_key()
209        result = 0
210        creds, inv = atest_gcp_utils.do_upload_flow(extra_args)
211        try:
212            verify_key = atest_utils.get_verify_key([test_infos[0].test_name],
213                                                    extra_args)
214            if extra_args.get(constants.VERIFY_ENV_VARIABLE, False):
215                # check environment variables.
216                atest_utils.handle_test_env_var(
217                    verify_key, result_path=constants.VERIFY_ENV_PATH)
218                return 0
219            # Change CWD to repo root to ensure TF can find prebuilt SDKs
220            # for some path-sensitive tests like robolectric.
221            os.chdir(os.path.abspath(os.getenv(constants.ANDROID_BUILD_TOP)))
222            if os.getenv(trb.OLD_OUTPUT_ENV_VAR):
223                result = self.run_tests_raw(test_infos, extra_args, reporter)
224            result = self.run_tests_pretty(test_infos, extra_args, reporter)
225        except atest_error.DryRunVerificationError as e:
226            atest_utils.colorful_print(str(e), constants.RED)
227            return ExitCode.VERIFY_FAILURE
228        finally:
229            if inv:
230                try:
231                    logging.disable(logging.INFO)
232                    # Always set invocation status to completed due to the ATest
233                    # handle whole process by its own.
234                    inv['schedulerState'] = 'completed'
235                    logstorage_utils.BuildClient(creds).update_invocation(inv)
236                    reporter.test_result_link = (constants.RESULT_LINK
237                                                 % inv['invocationId'])
238                finally:
239                    logging.disable(logging.NOTSET)
240        return result
241
242    def run_tests_raw(self, test_infos, extra_args, reporter):
243        """Run the list of test_infos. See base class for more.
244
245        Args:
246            test_infos: A list of TestInfos.
247            extra_args: Dict of extra args to add to test run.
248            reporter: An instance of result_report.ResultReporter.
249
250        Returns:
251            0 if tests succeed, non-zero otherwise.
252        """
253        iterations = self._generate_iterations(extra_args)
254        reporter.register_unsupported_runner(self.NAME)
255
256        ret_code = ExitCode.SUCCESS
257        for _ in range(iterations):
258            run_cmds = self.generate_run_commands(test_infos, extra_args)
259            subproc = self.run(run_cmds[0], output_to_stdout=True,
260                               env_vars=self.generate_env_vars(extra_args))
261            ret_code |= self.wait_for_subprocess(subproc)
262        return ret_code
263
264    def run_tests_pretty(self, test_infos, extra_args, reporter):
265        """Run the list of test_infos. See base class for more.
266
267        Args:
268            test_infos: A list of TestInfos.
269            extra_args: Dict of extra args to add to test run.
270            reporter: An instance of result_report.ResultReporter.
271
272        Returns:
273            0 if tests succeed, non-zero otherwise.
274        """
275        iterations = self._generate_iterations(extra_args)
276        ret_code = ExitCode.SUCCESS
277        for _ in range(iterations):
278            server = self._start_socket_server()
279            run_cmds = self.generate_run_commands(test_infos, extra_args,
280                                                  server.getsockname()[1])
281            subproc = self.run(run_cmds[0], output_to_stdout=self.is_verbose,
282                               env_vars=self.generate_env_vars(extra_args))
283            self.handle_subprocess(subproc, partial(self._start_monitor,
284                                                    server,
285                                                    subproc,
286                                                    reporter,
287                                                    extra_args))
288            server.close()
289            ret_code |= self.wait_for_subprocess(subproc)
290        return ret_code
291
292    # pylint: disable=too-many-branches
293    # pylint: disable=too-many-locals
294    def _start_monitor(self, server, tf_subproc, reporter, extra_args):
295        """Polling and process event.
296
297        Args:
298            server: Socket server object.
299            tf_subproc: The tradefed subprocess to poll.
300            reporter: Result_Reporter object.
301            extra_args: Dict of extra args to add to test run.
302        """
303        inputs = [server]
304        event_handlers = {}
305        data_map = {}
306        inv_socket = None
307        while inputs:
308            try:
309                readable, _, _ = select.select(inputs, [], [], SELECT_TIMEOUT)
310                for socket_object in readable:
311                    if socket_object is server:
312                        conn, addr = socket_object.accept()
313                        logging.debug('Accepted connection from %s', addr)
314                        conn.setblocking(False)
315                        inputs.append(conn)
316                        data_map[conn] = ''
317                        # The First connection should be invocation
318                        # level reporter.
319                        if not inv_socket:
320                            inv_socket = conn
321                    else:
322                        # Count invocation level reporter events
323                        # without showing real-time information.
324                        if inv_socket == socket_object:
325                            reporter.silent = True
326                            event_handler = event_handlers.setdefault(
327                                socket_object, EventHandler(reporter,
328                                                            self.NAME))
329                        else:
330                            event_handler = event_handlers.setdefault(
331                                socket_object, EventHandler(
332                                    result_reporter.ResultReporter(
333                                        collect_only=extra_args.get(
334                                            constants.COLLECT_TESTS_ONLY),
335                                        flakes_info=extra_args.get(
336                                            constants.FLAKES_INFO)),
337
338                                    self.NAME))
339                        recv_data = self._process_connection(data_map,
340                                                             socket_object,
341                                                             event_handler)
342                        if not recv_data:
343                            inputs.remove(socket_object)
344                            socket_object.close()
345            finally:
346                # Subprocess ended and all socket clients were closed.
347                if tf_subproc.poll() is not None and len(inputs) == 1:
348                    inputs.pop().close()
349                    if not reporter.all_test_results:
350                        atest_utils.colorful_print(
351                            r'No test to run. Test Logs have saved in '
352                            f'{reporter.log_path}.',
353                            constants.RED, constants.WHITE)
354                    if not data_map:
355                        metrics.LocalDetectEvent(
356                            detect_type=DetectType.TF_EXIT_CODE,
357                            result=tf_subproc.returncode)
358                        raise TradeFedExitError(tf_subproc.returncode)
359                    self._handle_log_associations(event_handlers)
360
361    def _process_connection(self, data_map, conn, event_handler):
362        """Process a socket connection betwen TF and ATest.
363
364        Expect data of form EVENT_NAME {JSON_DATA}.  Multiple events will be
365        \n deliminated.  Need to buffer data in case data exceeds socket
366        buffer.
367        E.q.
368            TEST_RUN_STARTED {runName":"hello_world_test","runAttempt":0}\n
369            TEST_STARTED {"start_time":2172917, "testName":"PrintHelloWorld"}\n
370        Args:
371            data_map: The data map of all connections.
372            conn: Socket connection.
373            event_handler: EventHandler object.
374
375        Returns:
376            True if conn.recv() has data , False otherwise.
377        """
378        # Set connection into blocking mode.
379        conn.settimeout(None)
380        data = conn.recv(SOCKET_BUFFER)
381        if isinstance(data, bytes):
382            data = data.decode()
383        logging.debug('received: %s', data)
384        if data:
385            data_map[conn] += data
386            while True:
387                match = EVENT_RE.match(data_map[conn])
388                if not match:
389                    break
390                try:
391                    event_data = json.loads(match.group('json_data'))
392                except ValueError:
393                    logging.debug('Json incomplete, wait for more data')
394                    break
395                event_name = match.group('event_name')
396                event_handler.process_event(event_name, event_data)
397                data_map[conn] = data_map[conn][match.end():]
398        return bool(data)
399
400    def _start_socket_server(self):
401        """Start a TCP server."""
402        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
403        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
404        # Port 0 lets the OS pick an open port between 1024 and 65535.
405        server.bind((SOCKET_HOST, 0))
406        server.listen(SOCKET_QUEUE_MAX)
407        server.settimeout(POLL_FREQ_SECS)
408        logging.debug('Socket server started on port %s',
409                      server.getsockname()[1])
410        return server
411
412    def generate_env_vars(self, extra_args):
413        """Convert extra args into env vars."""
414        env_vars = os.environ.copy()
415        if constants.TF_GLOBAL_CONFIG:
416            env_vars["TF_GLOBAL_CONFIG"] = constants.TF_GLOBAL_CONFIG
417        debug_port = extra_args.get(constants.TF_DEBUG, '')
418        if debug_port:
419            env_vars['TF_DEBUG'] = 'true'
420            env_vars['TF_DEBUG_PORT'] = str(debug_port)
421        filtered_paths = []
422        for path in str(env_vars['PYTHONPATH']).split(':'):
423            # TODO (b/166216843) Remove the hacky PYTHON path workaround.
424            if (str(path).startswith('/tmp/Soong.python_') and
425                    str(path).find('googleapiclient') > 0):
426                continue
427            filtered_paths.append(path)
428        env_vars['PYTHONPATH'] = ':'.join(filtered_paths)
429
430        # Use prebuilt aapt if there's no aapt under android system path which
431        # is aligned with build system.
432        # https://android.googlesource.com/platform/build/+/master/core/config.mk#529
433        if self._is_missing_exec(_AAPT):
434            prebuilt_aapt = Path.joinpath(
435                atest_utils.get_prebuilt_sdk_tools_dir(), _AAPT)
436            if os.path.exists(prebuilt_aapt):
437                env_vars['PATH'] = (str(prebuilt_aapt.parent) + ':'
438                                    + env_vars['PATH'])
439        return env_vars
440
441    # pylint: disable=unnecessary-pass
442    # Please keep above disable flag to ensure host_env_check is overriden.
443    def host_env_check(self):
444        """Check that host env has everything we need.
445
446        We actually can assume the host env is fine because we have the same
447        requirements that atest has. Update this to check for android env vars
448        if that changes.
449        """
450        pass
451
452    @staticmethod
453    def _is_missing_exec(executable):
454        """Check if system build executable is available.
455
456        Args:
457            executable: Executable we are checking for.
458
459        Returns:
460            True if executable is missing, False otherwise.
461        """
462        output = shutil.which(executable)
463        if not output:
464            return True
465        # TODO: Check if there is a clever way to determine if system adb is
466        # good enough.
467        root_dir = os.environ.get(constants.ANDROID_BUILD_TOP, '')
468        return os.path.commonprefix([output, root_dir]) != root_dir
469
470    def get_test_runner_build_reqs(self, test_infos: List[test_info.TestInfo]):
471        """Return the build requirements.
472
473        Args:
474            test_infos: List of TestInfo.
475
476        Returns:
477            Set of build targets.
478        """
479        build_req = self._BUILD_REQ.copy()
480        # Use different base build requirements if google-tf is around.
481        if self.module_info.is_module(constants.GTF_MODULE):
482            build_req = {constants.GTF_TARGET}
483        # Always add ATest's own TF target.
484        build_req.add(constants.ATEST_TF_MODULE)
485        # Add adb if we can't find it.
486        for executable in EXEC_DEPENDENCIES:
487            if self._is_missing_exec(executable):
488                if self.module_info.is_module(executable):
489                    build_req.add(executable)
490
491        # Force rebuilt all jars under $ANDROID_HOST_OUT to prevent old version
492        # host jars break the test.
493        build_req |= self._get_host_framework_targets()
494
495        build_req |= trb.gather_build_targets(test_infos)
496        return build_req
497
498    def _get_host_framework_targets(self) -> Set[str]:
499        """Get the build targets for all the existing jars under host framework.
500
501        Returns:
502            A set of build target name under $(ANDROID_HOST_OUT)/framework.
503        """
504        host_targets = set()
505        if not self.module_info:
506            return host_targets
507
508        framework_host_dir = Path(
509            os.environ.get(constants.ANDROID_HOST_OUT)).joinpath('framework')
510        if framework_host_dir.is_dir():
511            jars = framework_host_dir.glob('*.jar')
512            for jar in jars:
513                if self.module_info.is_module(jar.stem):
514                    host_targets.add(jar.stem)
515            logging.debug('Found exist host framework target:%s', host_targets)
516        return host_targets
517
518    def _parse_extra_args(self, test_infos, extra_args):
519        """Convert the extra args into something tf can understand.
520
521        Args:
522            extra_args: Dict of args
523
524        Returns:
525            Tuple of args to append and args not supported.
526        """
527        args_to_append, args_not_supported = extra_args_to_tf_args(
528            self.module_info, test_infos, extra_args)
529
530        # Set exclude instant app annotation for non-instant mode run.
531        if (constants.INSTANT not in extra_args and
532            self._has_instant_app_config(test_infos, self.module_info)):
533            args_to_append.append(constants.TF_TEST_ARG)
534            args_to_append.append(
535                '{tf_class}:{option_name}:{option_value}'.format(
536                    tf_class=constants.TF_AND_JUNIT_CLASS,
537                    option_name=constants.TF_EXCLUDE_ANNOTATE,
538                    option_value=constants.INSTANT_MODE_ANNOTATE))
539        # Force append --enable-parameterized-modules if args_to_append has
540        # --module-parameter in args_to_append
541        if constants.TF_MODULE_PARAMETER in args_to_append:
542            if constants.TF_ENABLE_PARAMETERIZED_MODULES not in args_to_append:
543                args_to_append.append(constants.TF_ENABLE_PARAMETERIZED_MODULES)
544        # If all the test config has config with auto enable parameter, force
545        # exclude those default parameters(ex: instant_app, secondary_user)
546        # TODO: (b/228433541) Remove the limitation after the root cause fixed.
547        if (len(test_infos) <= 1 and
548                self._is_all_tests_parameter_auto_enabled(test_infos)):
549            if constants.TF_ENABLE_PARAMETERIZED_MODULES not in args_to_append:
550                args_to_append.append(constants.TF_ENABLE_PARAMETERIZED_MODULES)
551                for exclude_parameter in constants.DEFAULT_EXCLUDE_PARAS:
552                    args_to_append.append('--exclude-module-parameters')
553                    args_to_append.append(exclude_parameter)
554        return args_to_append, args_not_supported
555
556    def _generate_metrics_folder(self, extra_args):
557        """Generate metrics folder."""
558        metrics_folder = ''
559        if extra_args.get(constants.PRE_PATCH_ITERATIONS):
560            metrics_folder = os.path.join(self.results_dir, 'baseline-metrics')
561        elif extra_args.get(constants.POST_PATCH_ITERATIONS):
562            metrics_folder = os.path.join(self.results_dir, 'new-metrics')
563        return metrics_folder
564
565    def _generate_iterations(self, extra_args):
566        """Generate iterations."""
567        iterations = 1
568        if extra_args.get(constants.PRE_PATCH_ITERATIONS):
569            iterations = extra_args.pop(constants.PRE_PATCH_ITERATIONS)
570        elif extra_args.get(constants.POST_PATCH_ITERATIONS):
571            iterations = extra_args.pop(constants.POST_PATCH_ITERATIONS)
572        return iterations
573
574    def generate_run_commands(self, test_infos, extra_args, port=None):
575        """Generate a single run command from TestInfos.
576
577        Args:
578            test_infos: A set of TestInfo instances.
579            extra_args: A Dict of extra args to append.
580            port: Optional. An int of the port number to send events to. If
581                  None, then subprocess reporter in TF won't try to connect.
582
583        Returns:
584            A list that contains the string of atest tradefed run command.
585            Only one command is returned.
586        """
587        args = self._create_test_args(test_infos)
588        metrics_folder = self._generate_metrics_folder(extra_args)
589
590        # Create a copy of args as more args could be added to the list.
591        test_args = list(args)
592        if port:
593            test_args.extend(['--subprocess-report-port', str(port)])
594        if metrics_folder:
595            test_args.extend(['--metrics-folder', metrics_folder])
596            logging.info('Saved metrics in: %s', metrics_folder)
597        if extra_args.get(constants.INVOCATION_ID, None):
598            test_args.append('--invocation-data invocation_id=%s'
599                             % extra_args[constants.INVOCATION_ID])
600        if extra_args.get(constants.WORKUNIT_ID, None):
601            test_args.append('--invocation-data work_unit_id=%s'
602                             % extra_args[constants.WORKUNIT_ID])
603        if extra_args.get(constants.LOCAL_BUILD_ID, None):
604            # TODO: (b/207584685) Replace with TF local build solutions.
605            test_args.append('--use-stub-build true')
606            test_args.append('--stub-build-id %s'
607                             % extra_args[constants.LOCAL_BUILD_ID])
608            test_args.append('--stub-build-target %s'
609                             % extra_args[constants.BUILD_TARGET])
610        if extra_args.get(constants.ENABLE_DEVICE_PREPARER, False):
611            test_args.append('--template:map preparers=%s'
612                             % constants.DEVICE_SETUP_PREPARER)
613        for info in test_infos:
614            if constants.TEST_WITH_MAINLINE_MODULES_RE.match(info.test_name):
615                # TODO(b/253641058) Remove this once mainline module
616                # binaries are stored under testcase directory.
617                self._copy_mainline_module_binary(info.mainline_modules)
618                test_args.append(constants.TF_ENABLE_MAINLINE_PARAMETERIZED_MODULES)
619                break
620        # For detailed logs, set TF options log-level/log-level-display as
621        # 'VERBOSE' by default.
622        log_level = 'VERBOSE'
623        test_args.extend(['--log-level-display', log_level])
624        test_args.extend(['--log-level', log_level])
625
626        # TODO(b/275110259) Remove this once TF not going to get bugreport.
627        test_args.extend(['--skip-all-system-status-check=true'])
628
629        # Set no-early-device-release by default to speed up TF teardown time.
630        if not constants.TF_EARLY_DEVICE_RELEASE in extra_args:
631            test_args.extend(['--no-early-device-release'])
632
633        args_to_add, args_not_supported = self._parse_extra_args(
634            test_infos, extra_args)
635
636        # If multiple devices in test config, automatically append
637        # --replicate-parent-setup and --multi-device-count
638        device_count = atest_configs.GLOBAL_ARGS.device_count_config
639        if device_count and device_count > 1:
640            args_to_add.append('--replicate-parent-setup')
641            args_to_add.append('--multi-device-count')
642            args_to_add.append(str(device_count))
643            os.environ.pop(constants.ANDROID_SERIAL, None)
644        else:
645            # TODO(b/122889707) Remove this after finding the root cause.
646            env_serial = os.environ.get(constants.ANDROID_SERIAL)
647            # Use the env variable ANDROID_SERIAL if it's set by user but only
648            # when the target tests are not deviceless tests.
649            if (env_serial and '--serial' not in args_to_add
650                and '-n' not in args_to_add):
651                args_to_add.append("--serial")
652                args_to_add.append(env_serial)
653
654        test_args.extend(args_to_add)
655        if args_not_supported:
656            logging.info('%s does not support the following args %s',
657                         self.EXECUTABLE, args_not_supported)
658
659        # Only need to check one TestInfo to determine if the tests are
660        # configured in TEST_MAPPING.
661        for_test_mapping = test_infos and test_infos[0].from_test_mapping
662        test_args.extend(atest_utils.get_result_server_args(for_test_mapping))
663        self.run_cmd_dict['args'] = ' '.join(test_args)
664        self.run_cmd_dict['tf_customize_template'] = (
665            self._extract_customize_tf_templates(extra_args, test_infos))
666
667        # Copy symbols if there are tests belong to native test.
668        self._handle_native_tests(test_infos)
669        return [self._RUN_CMD.format(**self.run_cmd_dict)]
670
671    def _flatten_test_infos(self, test_infos):
672        """Sort and group test_infos by module_name and sort and group filters
673        by class name.
674
675            Example of three test_infos in a set:
676                Module1, {(classA, {})}
677                Module1, {(classB, {Method1})}
678                Module1, {(classB, {Method2}}
679            Becomes a set with one element:
680                Module1, {(ClassA, {}), (ClassB, {Method1, Method2})}
681            Where:
682                  Each line is a test_info namedtuple
683                  {} = Frozenset
684                  () = TestFilter namedtuple
685
686        Args:
687            test_infos: A set of TestInfo namedtuples.
688
689        Returns:
690            A set of TestInfos flattened.
691        """
692        results = set()
693        key = lambda x: x.test_name
694        for module, group in atest_utils.sort_and_group(test_infos, key):
695            # module is a string, group is a generator of grouped TestInfos.
696            # Module Test, so flatten test_infos:
697            no_filters = False
698            filters = set()
699            test_runner = None
700            test_finder = None
701            build_targets = set()
702            data = {}
703            module_args = []
704            for test_info_i in group:
705                data.update(test_info_i.data)
706                # Extend data with constants.TI_MODULE_ARG instead of
707                # overwriting.
708                module_args.extend(test_info_i.data.get(
709                    constants.TI_MODULE_ARG, []))
710                test_runner = test_info_i.test_runner
711                test_finder = test_info_i.test_finder
712                build_targets |= test_info_i.build_targets
713                test_filters = test_info_i.data.get(constants.TI_FILTER)
714                if not test_filters or no_filters:
715                    # test_info wants whole module run, so hardcode no filters.
716                    no_filters = True
717                    filters = set()
718                    continue
719                filters |= test_filters
720            if module_args:
721                data[constants.TI_MODULE_ARG] = module_args
722            data[constants.TI_FILTER] = self._flatten_test_filters(filters)
723            results.add(
724                test_info.TestInfo(test_name=module,
725                                   test_runner=test_runner,
726                                   test_finder=test_finder,
727                                   build_targets=build_targets,
728                                   data=data))
729        return results
730
731    @staticmethod
732    def _flatten_test_filters(filters):
733        """Sort and group test_filters by class_name.
734
735            Example of three test_filters in a frozenset:
736                classA, {}
737                classB, {Method1}
738                classB, {Method2}
739            Becomes a frozenset with these elements:
740                classA, {}
741                classB, {Method1, Method2}
742            Where:
743                Each line is a TestFilter namedtuple
744                {} = Frozenset
745
746        Args:
747            filters: A frozenset of test_filters.
748
749        Returns:
750            A frozenset of test_filters flattened.
751        """
752        results = set()
753        key = lambda x: x.class_name
754        for class_name, group in atest_utils.sort_and_group(filters, key):
755            # class_name is a string, group is a generator of TestFilters
756            assert class_name is not None
757            methods = set()
758            for test_filter in group:
759                if not test_filter.methods:
760                    # Whole class should be run
761                    methods = set()
762                    break
763                methods |= test_filter.methods
764            results.add(test_info.TestFilter(class_name, frozenset(methods)))
765        return frozenset(results)
766
767    def _is_all_tests_parameter_auto_enabled(self, test_infos):
768        """Check if all the test infos are parameter auto enabled.
769
770        Args:
771            test_infos: A set of TestInfo instances.
772
773        Returns: True if all tests are parameter auto enabled, False otherwise.
774        """
775        for info in test_infos:
776            if not self._is_parameter_auto_enabled_cfg(info, self.module_info):
777                return False
778        return True
779
780    def _create_test_args(self, test_infos):
781        """Compile TF command line args based on the given test infos.
782
783        Args:
784            test_infos: A set of TestInfo instances.
785
786        Returns: A list of TF arguments to run the tests.
787        """
788        args = []
789        if not test_infos:
790            return []
791
792        test_infos = self._flatten_test_infos(test_infos)
793        has_integration_test = False
794
795        # Because current --include-filter arg will not working if ATest pass
796        # both --module and --include-filter to TF, only test by --module will
797        # be run. Make a check first, only use --module if all tests are all
798        # parameter auto enabled.
799        # Only auto-enable the parameter if there's only one test.
800        # TODO: (b/228433541) Remove the limitation after the root cause fixed.
801        use_module_arg = False
802        if len(test_infos) <= 1:
803            use_module_arg = self._is_all_tests_parameter_auto_enabled(
804                test_infos)
805
806        for info in test_infos:
807            # Integration test exists in TF's jar, so it must have the option
808            # if it's integration finder.
809            if info.test_finder in _INTEGRATION_FINDERS:
810                has_integration_test = True
811            # For non-paramertize test module, use --include-filter, but for
812            # tests which have auto enable paramertize config use --module
813            # instead.
814            if (use_module_arg
815                and self._is_parameter_auto_enabled_cfg(
816                    info, self.module_info)):
817                args.extend([constants.TF_MODULE_FILTER, info.test_name])
818            else:
819                args.extend([constants.TF_INCLUDE_FILTER, info.test_name])
820            for option in info.data.get(constants.TI_MODULE_ARG, []):
821                if constants.TF_INCLUDE_FILTER_OPTION == option[0]:
822                    suite_filter = (
823                        constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
824                            test_name=info.test_name, option_value=option[1]))
825                    args.extend([constants.TF_INCLUDE_FILTER, suite_filter])
826                elif constants.TF_EXCLUDE_FILTER_OPTION == option[0]:
827                    suite_filter = (
828                        constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
829                            test_name=info.test_name, option_value=option[1]))
830                    args.extend([constants.TF_EXCLUDE_FILTER, suite_filter])
831                else:
832                    module_arg = (
833                        constants.TF_MODULE_ARG_VALUE_FMT.format(
834                            test_name=info.test_name, option_name=option[0],
835                            option_value=option[1]))
836                    args.extend([constants.TF_MODULE_ARG, module_arg])
837
838        # Add ATest include filter
839        args.extend(get_include_filter(test_infos))
840
841        # TODO (b/141090547) Pass the config path to TF to load configs.
842        # Compile option in TF if finder is not INTEGRATION or not set.
843        if not has_integration_test:
844            args.append(constants.TF_SKIP_LOADING_CONFIG_JAR)
845        return args
846
847    def _extract_rerun_options(self, extra_args):
848        """Extract rerun options to a string for output.
849
850        Args:
851            extra_args: Dict of extra args for test runners to use.
852
853        Returns: A string of rerun options.
854        """
855        extracted_options = ['{} {}'.format(arg, extra_args[arg])
856                             for arg in extra_args
857                             if arg in self._RERUN_OPTION_GROUP]
858        return ' '.join(extracted_options)
859
860    def _extract_customize_tf_templates(self, extra_args, test_infos):
861        """Extract tradefed template options to a string for output.
862
863        Args:
864            extra_args: Dict of extra args for test runners to use.
865            test_infos: A set of TestInfo instances.
866
867        Returns: A string of tradefed template options.
868        """
869        tf_templates = extra_args.get(constants.TF_TEMPLATE, [])
870        tf_template_keys = [i.split('=')[0] for i in tf_templates]
871        for info in test_infos:
872            if (info.aggregate_metrics_result
873                    and 'metric_post_processor' not in tf_template_keys):
874                template_key = 'metric_post_processor'
875                template_value = (
876                    'google/template/postprocessors/metric-file-aggregate')
877                tf_templates.append(f'{template_key}={template_value}')
878        return ' '.join(['--template:map %s' % x for x in tf_templates])
879
880    def _handle_log_associations(self, event_handlers):
881        """Handle TF's log associations information data.
882
883        log_association dict:
884        {'loggedFile': '/tmp/serial-util11375755456514097276.ser',
885         'dataName': 'device_logcat_setup_127.0.0.1:58331',
886         'time': 1602038599.856113},
887
888        Args:
889            event_handlers: Dict of {socket_object:EventHandler}.
890
891        """
892        log_associations = []
893        for _, event_handler in event_handlers.items():
894            if event_handler.log_associations:
895                log_associations += event_handler.log_associations
896        device_test_end_log_time = ''
897        device_teardown_log_time = ''
898        for log_association in log_associations:
899            if 'device_logcat_test' in log_association.get('dataName', ''):
900                device_test_end_log_time = log_association.get('time')
901            if 'device_logcat_teardown' in log_association.get('dataName', ''):
902                device_teardown_log_time = log_association.get('time')
903        if device_test_end_log_time and device_teardown_log_time:
904            teardowntime = (float(device_teardown_log_time) -
905                            float(device_test_end_log_time))
906            logging.debug('TF logcat teardown time=%s seconds.', teardowntime)
907            metrics.LocalDetectEvent(
908                detect_type=DetectType.TF_TEARDOWN_LOGCAT,
909                result=int(teardowntime))
910
911    @staticmethod
912    def _has_instant_app_config(test_infos, mod_info):
913        """Check if one of the input tests defined instant app mode in config.
914
915        Args:
916            test_infos: A set of TestInfo instances.
917            mod_info: ModuleInfo object.
918
919        Returns: True if one of the tests set up instant app mode.
920        """
921        for tinfo in test_infos:
922            test_config, _ = test_finder_utils.get_test_config_and_srcs(
923                tinfo, mod_info)
924            if test_config:
925                parameters = atest_utils.get_config_parameter(test_config)
926                if constants.TF_PARA_INSTANT_APP in parameters:
927                    return True
928        return False
929
930    @staticmethod
931    def _is_parameter_auto_enabled_cfg(tinfo, mod_info):
932        """Check if input tests contains auto enable support parameters.
933
934        Args:
935            test_infos: A set of TestInfo instances.
936            mod_info: ModuleInfo object.
937
938        Returns: True if input test has parameter setting which is not in the
939                 exclude list.
940        """
941        test_config, _ = test_finder_utils.get_test_config_and_srcs(
942            tinfo, mod_info)
943        if test_config:
944            parameters = atest_utils.get_config_parameter(test_config)
945            if (parameters - constants.DEFAULT_EXCLUDE_PARAS
946                - constants.DEFAULT_EXCLUDE_NOT_PARAS):
947                return True
948        return False
949
950    def _handle_native_tests(self, test_infos):
951        """Handling some extra tasks for running native tests from tradefed.
952
953        Args:
954            test_infos: A set of TestInfo instances.
955        """
956        for tinfo in test_infos:
957            test_config, _ = test_finder_utils.get_test_config_and_srcs(
958                tinfo, self.module_info)
959            if test_config:
960                module_name, device_path = atest_utils.get_config_gtest_args(
961                    test_config)
962                if module_name and device_path:
963                    atest_utils.copy_native_symbols(module_name, device_path)
964
965    # TODO(b/253641058) remove copying files once mainline module
966    # binaries are stored under testcase directory.
967    def _copy_mainline_module_binary(self, mainline_modules):
968        """Copies mainline module binaries to out/dist/mainline_modules_{arch}
969
970        Copies the mainline module binaries to the location that
971        MainlineModuleHandler in TF expects since there is no way to
972        explicitly tweak the search path.
973
974        Args:
975            mainline_modules: A list of mainline modules.
976        """
977        config = atest_utils.get_android_config()
978        arch = config.get('TARGET_ARCH')
979        dest_dir = atest_utils.DIST_OUT_DIR.joinpath(f'mainline_modules_{arch}')
980        dest_dir.mkdir(parents=True, exist_ok=True)
981
982        for module in mainline_modules:
983            target_module_info = self.module_info.get_module_info(module)
984            installed_paths = target_module_info[constants.MODULE_INSTALLED]
985
986            for installed_path in installed_paths:
987                if not re.search(atest_utils.MAINLINE_MODULES_EXT_RE, installed_path):
988                    atest_utils.colorful_print(
989                        '%s is not a apk or apex file. Did you run mainline '
990                        'local setup script? Please refer to %s' %
991                        (installed_path, MAINLINE_LOCAL_DOC),
992                        constants.YELLOW)
993                    continue
994                file_name = Path(installed_path).name
995                dest_path = Path(dest_dir).joinpath(file_name)
996                if dest_path.exists():
997                    atest_utils.colorful_print(
998                        'Replacing APEX in %s with %s' % (dest_path, installed_path),
999                        constants.CYAN)
1000                    logging.debug(
1001                        'deleting the old file: %s and copy a new binary',
1002                        dest_path)
1003                    dest_path.unlink()
1004                shutil.copyfile(installed_path, dest_path)
1005
1006                break
1007
1008def generate_annotation_filter_args(
1009        arg_value: Any, mod_info: module_info.ModuleInfo,
1010        test_infos: List[test_info.TestInfo]) -> List[str]:
1011    """Generate TF annotation filter arguments.
1012
1013    Args:
1014        arg_value: Argument value for annotation filter.
1015        mod_info: ModuleInfo object.
1016        test_infos: A set of TestInfo instances.
1017
1018    Returns:
1019        List of TF annotation filter arguments.
1020    """
1021    annotation_filter_args = []
1022    for info in test_infos:
1023        test_name = info.test_name
1024        for keyword in arg_value:
1025            annotation = atest_utils.get_full_annotation_class_name(
1026                mod_info.get_module_info(test_name), keyword)
1027            if annotation:
1028                module_arg = (constants.TF_MODULE_ARG_VALUE_FMT.format(
1029                    test_name=test_name,
1030                    option_name=constants.INCLUDE_ANNOTATION,
1031                    option_value=annotation))
1032                annotation_filter_args.extend([constants.TF_MODULE_ARG, module_arg])
1033            logging.error(
1034                atest_utils.colorize(
1035                    f'Cannot find similar annotation: {keyword}',
1036                    constants.RED))
1037    return annotation_filter_args
1038
1039
1040def extra_args_to_tf_args(
1041    mod_info: module_info.ModuleInfo,
1042    test_infos: List[test_info.TestInfo],
1043    extra_args: Dict[str, Any],
1044) -> Tuple[Dict[str, Any], Dict[str, Any]]:
1045    """Convert the extra args into atest_tf_test_runner supported args.
1046
1047    Args:
1048        mod_info: ModuleInfo object.
1049        test_infos: A set of TestInfo instances.
1050        extra_args: Dict of args
1051
1052    Returns:
1053        Tuple of ARGS that atest_tf supported and not supported.
1054    """
1055    supported_args = []
1056    unsupported_args = []
1057
1058    def constant_list(*value):
1059        return lambda *_: value
1060
1061    # pylint: disable=unused-argument
1062    def print_message(message):
1063        def inner(*args):
1064            print(message)
1065            return []
1066        return inner
1067
1068    # Mapping supported TF arguments to the processing function.
1069    supported_tf_args = dict({
1070        constants.WAIT_FOR_DEBUGGER:
1071            constant_list('--wait-for-debugger'),
1072        constants.DISABLE_INSTALL:
1073            constant_list('--disable-target-preparers'),
1074        constants.SERIAL:
1075            lambda arg_value, *_:
1076            [j for d in arg_value for j in ('--serial', d)],
1077        constants.SHARDING:
1078            lambda arg_value, *_: ['--shard-count',
1079                                   str(arg_value)],
1080        constants.DISABLE_TEARDOWN:
1081            constant_list('--disable-teardown'),
1082        constants.HOST:
1083            constant_list('-n', '--prioritize-host-config',
1084                          '--skip-host-arch-check'),
1085        constants.CUSTOM_ARGS:
1086            # custom args value is a list.
1087            lambda arg_value, *_: arg_value,
1088        constants.ALL_ABI:
1089            constant_list('--all-abi'),
1090        constants.INSTANT:
1091            constant_list(constants.TF_ENABLE_PARAMETERIZED_MODULES,
1092                          constants.TF_MODULE_PARAMETER, 'instant_app'),
1093        constants.USER_TYPE:
1094            lambda arg_value, *_: [
1095                constants.TF_ENABLE_PARAMETERIZED_MODULES,
1096                '--enable-optional-parameterization',
1097                constants.TF_MODULE_PARAMETER,
1098                str(arg_value)
1099            ],
1100        constants.ITERATIONS:
1101            lambda arg_value, *_: [
1102                '--retry-strategy', constants.ITERATIONS,
1103                '--max-testcase-run-count', str(arg_value)
1104            ],
1105        constants.RERUN_UNTIL_FAILURE:
1106            lambda arg_value, *_: [
1107                '--retry-strategy', constants.RERUN_UNTIL_FAILURE,
1108                '--max-testcase-run-count', str(arg_value)
1109            ],
1110        constants.RETRY_ANY_FAILURE:
1111            lambda arg_value, *_: [
1112                '--retry-strategy', constants.RETRY_ANY_FAILURE,
1113                '--max-testcase-run-count', str(arg_value)
1114            ],
1115        constants.COLLECT_TESTS_ONLY:
1116            constant_list('--collect-tests-only'),
1117        constants.NO_ENABLE_ROOT:
1118            constant_list('--no-enable-root'),
1119        constants.TF_DEBUG:
1120            print_message("Please attach process to your IDE..."),
1121        constants.ANNOTATION_FILTER:
1122            generate_annotation_filter_args,
1123        constants.TEST_FILTER:
1124            lambda arg_value, *_: [
1125                '--test-arg',
1126                'com.android.tradefed.testtype.AndroidJUnitTest:'
1127                f'include-filter:{arg_value}',
1128                '--test-arg',
1129                'com.android.tradefed.testtype.GTest:native-test-flag:'
1130                f'--gtest_filter={arg_value}',
1131                '--test-arg',
1132                'com.android.tradefed.testtype.HostGTest:native-test-flag:'
1133                f'--gtest_filter={arg_value}'
1134            ],
1135        constants.TEST_TIMEOUT:
1136            lambda arg_value, *_: [
1137                '--test-arg',
1138                'com.android.tradefed.testtype.AndroidJUnitTest:'
1139                f'shell-timeout:{arg_value}',
1140                '--test-arg',
1141                'com.android.tradefed.testtype.AndroidJUnitTest:'
1142                f'test-timeout:{arg_value}',
1143                '--test-arg',
1144                'com.android.tradefed.testtype.HostGTest:'
1145                f'native-test-timeout:{arg_value}',
1146                '--test-arg',
1147                'com.android.tradefed.testtype.GTest:'
1148                f'native-test-timeout:{arg_value}',
1149            ],
1150        constants.COVERAGE: coverage.tf_args,
1151    })
1152
1153    for arg in extra_args:
1154        if arg in supported_tf_args:
1155            tf_args = supported_tf_args[arg](extra_args[arg], mod_info,
1156                                             test_infos)
1157            if tf_args:
1158                supported_args.extend(tf_args)
1159            continue
1160
1161        if arg in (constants.TF_TEMPLATE,
1162                   constants.TF_EARLY_DEVICE_RELEASE,
1163                   constants.INVOCATION_ID,
1164                   constants.WORKUNIT_ID,
1165                   constants.REQUEST_UPLOAD_RESULT,
1166                   constants.DISABLE_UPLOAD_RESULT,
1167                   constants.LOCAL_BUILD_ID,
1168                   constants.BUILD_TARGET,
1169                   constants.ENABLE_DEVICE_PREPARER,
1170                   constants.DRY_RUN,
1171                   constants.VERIFY_ENV_VARIABLE,
1172                   constants.FLAKES_INFO,
1173                   constants.LD_LIBRARY_PATH,
1174                   constants.DEVICE_ONLY):
1175            continue
1176        unsupported_args.append(arg)
1177    return supported_args, unsupported_args
1178
1179def get_include_filter(test_infos: List[test_info.TestInfo]) -> List[str]:
1180    """Generate a list of tradefed filter argument from TestInfos.
1181
1182    The tradefed argument format should be:
1183    atest-include-filter <module-name>:<include-filter-value>
1184    """
1185    tf_args = []
1186    for info in test_infos:
1187        filters = set()
1188        for test_info_filter in info.data.get(constants.TI_FILTER, []):
1189            filters.update(test_info_filter.to_set_of_tf_strings())
1190        for test_filter in filters:
1191            filter_arg = constants.TF_ATEST_INCLUDE_FILTER_VALUE_FMT.format(
1192                test_name=info.test_name, test_filter=test_filter)
1193            tf_args.extend([constants.TF_ATEST_INCLUDE_FILTER, filter_arg])
1194    return tf_args
1195