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
19from __future__ import print_function
20
21import json
22import logging
23import os
24import re
25import select
26import shutil
27import socket
28import uuid
29
30from functools import partial
31from pathlib import Path
32
33import atest_utils
34import constants
35import result_reporter
36
37from logstorage import atest_gcp_utils
38from logstorage import logstorage_utils
39from metrics import metrics
40from test_finders import test_finder_utils
41from test_finders import test_info
42from test_runners import test_runner_base
43from .event_handler import EventHandler
44
45POLL_FREQ_SECS = 10
46SOCKET_HOST = '127.0.0.1'
47SOCKET_QUEUE_MAX = 1
48SOCKET_BUFFER = 4096
49SELECT_TIMEOUT = 0.5
50
51# Socket Events of form FIRST_EVENT {JSON_DATA}\nSECOND_EVENT {JSON_DATA}
52# EVENT_RE has groups for the name and the data. "." does not match \n.
53EVENT_RE = re.compile(r'\n*(?P<event_name>[A-Z_]+) (?P<json_data>{.*})(?=\n|.)*')
54
55EXEC_DEPENDENCIES = ('adb', 'aapt', 'fastboot')
56
57TRADEFED_EXIT_MSG = 'TradeFed subprocess exited early with exit code=%s.'
58
59LOG_FOLDER_NAME = 'log'
60
61_INTEGRATION_FINDERS = frozenset(['', 'INTEGRATION', 'INTEGRATION_FILE_PATH'])
62
63class TradeFedExitError(Exception):
64    """Raised when TradeFed exists before test run has finished."""
65
66
67class AtestTradefedTestRunner(test_runner_base.TestRunnerBase):
68    """TradeFed Test Runner class."""
69    NAME = 'AtestTradefedTestRunner'
70    EXECUTABLE = 'atest_tradefed.sh'
71    _TF_TEMPLATE = 'template/atest_local_min'
72    # Use --no-enable-granular-attempts to control reporter replay behavior.
73    # TODO(b/142630648): Enable option enable-granular-attempts
74    # in sharding mode.
75    _LOG_ARGS = ('--logcat-on-failure --atest-log-file-path={log_path} '
76                 '--no-enable-granular-attempts '
77                 '--proto-output-file={proto_path}')
78    _RUN_CMD = ('{env} {exe} {template} --template:map '
79                'test=atest {tf_customize_template} {log_args} {args}')
80    _BUILD_REQ = {'tradefed-core'}
81    _RERUN_OPTION_GROUP = [constants.ITERATIONS,
82                           constants.RERUN_UNTIL_FAILURE,
83                           constants.RETRY_ANY_FAILURE]
84
85    def __init__(self, results_dir, module_info=None, **kwargs):
86        """Init stuff for base class."""
87        super().__init__(results_dir, **kwargs)
88        self.module_info = module_info
89        self.log_path = os.path.join(results_dir, LOG_FOLDER_NAME)
90        if not os.path.exists(self.log_path):
91            os.makedirs(self.log_path)
92        log_args = {'log_path': self.log_path,
93                    'proto_path': os.path.join(self.results_dir, constants.ATEST_TEST_RECORD_PROTO)}
94        self.run_cmd_dict = {'env': self._get_ld_library_path(),
95                             'exe': self.EXECUTABLE,
96                             'template': self._TF_TEMPLATE,
97                             'tf_customize_template': '',
98                             'args': '',
99                             'log_args': self._LOG_ARGS.format(**log_args)}
100        self.is_verbose = logging.getLogger().isEnabledFor(logging.DEBUG)
101        self.root_dir = os.environ.get(constants.ANDROID_BUILD_TOP)
102
103    def _get_ld_library_path(self):
104        """Get the extra environment setup string for running TF.
105
106        Returns:
107            Strings for the environment passed to TF. Currently only
108            LD_LIBRARY_PATH for TF to load the correct local shared libraries.
109        """
110        out_dir = os.environ.get(constants.ANDROID_HOST_OUT, '')
111        lib_dirs = ['lib', 'lib64']
112        path = ''
113        for lib in lib_dirs:
114            lib_dir = os.path.join(out_dir, lib)
115            path = path + lib_dir + ':'
116        return 'LD_LIBRARY_PATH=%s' % path
117
118    def _try_set_gts_authentication_key(self):
119        """Set GTS authentication key if it is available or exists.
120
121        Strategy:
122            Get APE_API_KEY from os.environ:
123                - If APE_API_KEY is already set by user -> do nothing.
124            Get the APE_API_KEY from constants:
125                - If the key file exists -> set to env var.
126            If APE_API_KEY isn't set and the key file doesn't exist:
127                - Warn user some GTS tests may fail without authentication.
128        """
129        if os.environ.get('APE_API_KEY'):
130            logging.debug('APE_API_KEY is set by developer.')
131            return
132        ape_api_key = constants.GTS_GOOGLE_SERVICE_ACCOUNT
133        key_path = os.path.join(self.root_dir, ape_api_key)
134        if ape_api_key and os.path.exists(key_path):
135            logging.debug('Set APE_API_KEY: %s', ape_api_key)
136            os.environ['APE_API_KEY'] = key_path
137        else:
138            logging.debug('APE_API_KEY not set, some GTS tests may fail'
139                          ' without authentication.')
140
141    def run_tests(self, test_infos, extra_args, reporter):
142        """Run the list of test_infos. See base class for more.
143
144        Args:
145            test_infos: A list of TestInfos.
146            extra_args: Dict of extra args to add to test run.
147            reporter: An instance of result_report.ResultReporter.
148
149        Returns:
150            0 if tests succeed, non-zero otherwise.
151        """
152        reporter.log_path = self.log_path
153        reporter.rerun_options = self._extract_rerun_options(extra_args)
154        # Set google service key if it's available or found before
155        # running tests.
156        self._try_set_gts_authentication_key()
157        result = 0
158        creds, inv = self._do_upload_flow(extra_args)
159        try:
160            if os.getenv(test_runner_base.OLD_OUTPUT_ENV_VAR):
161                result = self.run_tests_raw(test_infos, extra_args, reporter)
162            result = self.run_tests_pretty(test_infos, extra_args, reporter)
163        finally:
164            if inv:
165                try:
166                    logging.disable(logging.INFO)
167                    # Always set invocation status to completed due to the ATest
168                    # handle whole process by its own.
169                    inv['schedulerState'] = 'completed'
170                    logstorage_utils.BuildClient(creds).update_invocation(inv)
171                    reporter.test_result_link = (constants.RESULT_LINK
172                                                 % inv['invocationId'])
173                finally:
174                    logging.disable(logging.NOTSET)
175        return result
176
177    def _do_upload_flow(self, extra_args):
178        """Run upload flow.
179
180        Asking user's decision and do the related steps.
181
182        Args:
183            extra_args: Dict of extra args to add to test run.
184        Return:
185            tuple(invocation, workunit)
186        """
187        config_folder = os.path.join(os.path.expanduser('~'), '.atest')
188        creds = self._request_consent_of_upload_test_result(
189            config_folder,
190            extra_args.get(constants.REQUEST_UPLOAD_RESULT, None))
191        if creds:
192            inv, workunit = self._prepare_data(creds)
193            extra_args[constants.INVOCATION_ID] = inv['invocationId']
194            extra_args[constants.WORKUNIT_ID] = workunit['id']
195            if not os.path.exists(os.path.dirname(constants.TOKEN_FILE_PATH)):
196                os.makedirs(os.path.dirname(constants.TOKEN_FILE_PATH))
197            with open(constants.TOKEN_FILE_PATH, 'w') as token_file:
198                token_file.write(creds.token_response['access_token'])
199            return creds, inv
200        return None, None
201
202    def _prepare_data(self, creds):
203        """Prepare data for build api using.
204
205        Args:
206            creds: The credential object.
207        Return:
208            invocation and workunit object.
209        """
210        try:
211            logging.disable(logging.INFO)
212            external_id = str(uuid.uuid4())
213            client = logstorage_utils.BuildClient(creds)
214            branch = self._get_branch(client)
215            target = self._get_target(branch, client)
216            build_record = client.insert_local_build(external_id,
217                                                     target,
218                                                     branch)
219            client.insert_build_attempts(build_record)
220            invocation = client.insert_invocation(build_record)
221            workunit = client.insert_work_unit(invocation)
222            return invocation, workunit
223        finally:
224            logging.disable(logging.NOTSET)
225
226    def _get_branch(self, build_client):
227        """Get source code tree branch.
228
229        Args:
230            build_client: The build client object.
231        Return:
232            "git_master" in internal git, "aosp-master" otherwise.
233        """
234        default_branch = ('git_master'
235                          if constants.CREDENTIAL_FILE_NAME else 'aosp-master')
236        local_branch = atest_utils.get_manifest_branch()
237        branches = [b['name'] for b in build_client.list_branch()['branches']]
238        return local_branch if local_branch in branches else default_branch
239
240    def _get_target(self, branch, build_client):
241        """Get local build selected target.
242
243        Args:
244            branch: The branch want to check.
245            build_client: The build client object.
246        Return:
247            The matched build target, "aosp_x86-userdebug" otherwise.
248        """
249        default_target = 'aosp_x86-userdebug'
250        local_target = atest_utils.get_build_target()
251        targets = [t['target']
252                   for t in build_client.list_target(branch)['targets']]
253        return local_target if local_target in targets else default_target
254
255    def _request_consent_of_upload_test_result(self, config_folder,
256                                               request_to_upload_result):
257        """Request the consent of upload test results at the first time.
258
259        Args:
260            config_folder: The directory path to put config file.
261            request_to_upload_result: Prompt message for user determine.
262        Return:
263            The credential object.
264        """
265        if not os.path.exists(config_folder):
266            os.makedirs(config_folder)
267        not_upload_file = os.path.join(config_folder,
268                                       constants.DO_NOT_UPLOAD)
269        # Do nothing if there are no related config or DO_NOT_UPLOAD exists.
270        if (not constants.CREDENTIAL_FILE_NAME or
271                not constants.TOKEN_FILE_PATH):
272            return None
273
274        creds_f = os.path.join(config_folder, constants.CREDENTIAL_FILE_NAME)
275        if request_to_upload_result:
276            if os.path.exists(not_upload_file):
277                os.remove(not_upload_file)
278            if os.path.exists(creds_f):
279                os.remove(creds_f)
280
281        # If the credential file exists or the user says “Yes”, ATest will
282        # try to get the credential from the file, else will create a
283        # DO_NOT_UPLOAD to keep the user's decision.
284        if not os.path.exists(not_upload_file):
285            if (os.path.exists(creds_f) or
286                    (request_to_upload_result and
287                     atest_utils.prompt_with_yn_result(
288                         constants.UPLOAD_TEST_RESULT_MSG, False))):
289                return atest_gcp_utils.GCPHelper(
290                    client_id=constants.CLIENT_ID,
291                    client_secret=constants.CLIENT_SECRET,
292                    user_agent='atest').get_credential_with_auth_flow(creds_f)
293
294        Path(not_upload_file).touch()
295        return None
296
297    def run_tests_raw(self, test_infos, extra_args, reporter):
298        """Run the list of test_infos. See base class for more.
299
300        Args:
301            test_infos: A list of TestInfos.
302            extra_args: Dict of extra args to add to test run.
303            reporter: An instance of result_report.ResultReporter.
304
305        Returns:
306            0 if tests succeed, non-zero otherwise.
307        """
308        iterations = self._generate_iterations(extra_args)
309        reporter.register_unsupported_runner(self.NAME)
310
311        ret_code = constants.EXIT_CODE_SUCCESS
312        for _ in range(iterations):
313            run_cmds = self.generate_run_commands(test_infos, extra_args)
314            subproc = self.run(run_cmds[0], output_to_stdout=True,
315                               env_vars=self.generate_env_vars(extra_args))
316            ret_code |= self.wait_for_subprocess(subproc)
317        return ret_code
318
319    def run_tests_pretty(self, test_infos, extra_args, reporter):
320        """Run the list of test_infos. See base class for more.
321
322        Args:
323            test_infos: A list of TestInfos.
324            extra_args: Dict of extra args to add to test run.
325            reporter: An instance of result_report.ResultReporter.
326
327        Returns:
328            0 if tests succeed, non-zero otherwise.
329        """
330        iterations = self._generate_iterations(extra_args)
331        ret_code = constants.EXIT_CODE_SUCCESS
332        for _ in range(iterations):
333            server = self._start_socket_server()
334            run_cmds = self.generate_run_commands(test_infos, extra_args,
335                                                  server.getsockname()[1])
336            subproc = self.run(run_cmds[0], output_to_stdout=self.is_verbose,
337                               env_vars=self.generate_env_vars(extra_args))
338            self.handle_subprocess(subproc, partial(self._start_monitor,
339                                                    server,
340                                                    subproc,
341                                                    reporter,
342                                                    extra_args))
343            server.close()
344            ret_code |= self.wait_for_subprocess(subproc)
345        return ret_code
346
347    # pylint: disable=too-many-branches
348    # pylint: disable=too-many-locals
349    def _start_monitor(self, server, tf_subproc, reporter, extra_args):
350        """Polling and process event.
351
352        Args:
353            server: Socket server object.
354            tf_subproc: The tradefed subprocess to poll.
355            reporter: Result_Reporter object.
356            extra_args: Dict of extra args to add to test run.
357        """
358        inputs = [server]
359        event_handlers = {}
360        data_map = {}
361        inv_socket = None
362        while inputs:
363            try:
364                readable, _, _ = select.select(inputs, [], [], SELECT_TIMEOUT)
365                for socket_object in readable:
366                    if socket_object is server:
367                        conn, addr = socket_object.accept()
368                        logging.debug('Accepted connection from %s', addr)
369                        conn.setblocking(False)
370                        inputs.append(conn)
371                        data_map[conn] = ''
372                        # The First connection should be invocation
373                        # level reporter.
374                        if not inv_socket:
375                            inv_socket = conn
376                    else:
377                        # Count invocation level reporter events
378                        # without showing real-time information.
379                        if inv_socket == socket_object:
380                            reporter.silent = True
381                            event_handler = event_handlers.setdefault(
382                                socket_object, EventHandler(reporter,
383                                                            self.NAME))
384                        else:
385                            event_handler = event_handlers.setdefault(
386                                socket_object, EventHandler(
387                                    result_reporter.ResultReporter(
388                                        collect_only=extra_args.get(
389                                            constants.COLLECT_TESTS_ONLY),
390                                        flakes_info=extra_args.get(
391                                            constants.FLAKES_INFO)),
392
393                                    self.NAME))
394                        recv_data = self._process_connection(data_map,
395                                                             socket_object,
396                                                             event_handler)
397                        if not recv_data:
398                            inputs.remove(socket_object)
399                            socket_object.close()
400            finally:
401                # Subprocess ended and all socket clients were closed.
402                if tf_subproc.poll() is not None and len(inputs) == 1:
403                    inputs.pop().close()
404                    if not reporter.all_test_results:
405                        atest_utils.colorful_print(
406                            r'No test to run. Please check: '
407                            r'{} for detail.'.format(reporter.log_path),
408                            constants.RED, highlight=True)
409                    if not data_map:
410                        raise TradeFedExitError(TRADEFED_EXIT_MSG
411                                                % tf_subproc.returncode)
412                    self._handle_log_associations(event_handlers)
413
414    def _process_connection(self, data_map, conn, event_handler):
415        """Process a socket connection betwen TF and ATest.
416
417        Expect data of form EVENT_NAME {JSON_DATA}.  Multiple events will be
418        \n deliminated.  Need to buffer data in case data exceeds socket
419        buffer.
420        E.q.
421            TEST_RUN_STARTED {runName":"hello_world_test","runAttempt":0}\n
422            TEST_STARTED {"start_time":2172917, "testName":"PrintHelloWorld"}\n
423        Args:
424            data_map: The data map of all connections.
425            conn: Socket connection.
426            event_handler: EventHandler object.
427
428        Returns:
429            True if conn.recv() has data , False otherwise.
430        """
431        # Set connection into blocking mode.
432        conn.settimeout(None)
433        data = conn.recv(SOCKET_BUFFER)
434        if isinstance(data, bytes):
435            data = data.decode()
436        logging.debug('received: %s', data)
437        if data:
438            data_map[conn] += data
439            while True:
440                match = EVENT_RE.match(data_map[conn])
441                if not match:
442                    break
443                try:
444                    event_data = json.loads(match.group('json_data'))
445                except ValueError:
446                    logging.debug('Json incomplete, wait for more data')
447                    break
448                event_name = match.group('event_name')
449                event_handler.process_event(event_name, event_data)
450                data_map[conn] = data_map[conn][match.end():]
451        return bool(data)
452
453    def _start_socket_server(self):
454        """Start a TCP server."""
455        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
456        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
457        # Port 0 lets the OS pick an open port between 1024 and 65535.
458        server.bind((SOCKET_HOST, 0))
459        server.listen(SOCKET_QUEUE_MAX)
460        server.settimeout(POLL_FREQ_SECS)
461        logging.debug('Socket server started on port %s',
462                      server.getsockname()[1])
463        return server
464
465    def generate_env_vars(self, extra_args):
466        """Convert extra args into env vars."""
467        env_vars = os.environ.copy()
468        if constants.TF_GLOBAL_CONFIG:
469            env_vars["TF_GLOBAL_CONFIG"] = constants.TF_GLOBAL_CONFIG
470        debug_port = extra_args.get(constants.TF_DEBUG, '')
471        if debug_port:
472            env_vars['TF_DEBUG'] = 'true'
473            env_vars['TF_DEBUG_PORT'] = str(debug_port)
474        filtered_paths = []
475        for path in str(env_vars['PYTHONPATH']).split(':'):
476            # TODO (b/166216843) Remove the hacky PYTHON path workaround.
477            if (str(path).startswith('/tmp/Soong.python_') and
478                    str(path).find('googleapiclient') > 0):
479                continue
480            filtered_paths.append(path)
481        env_vars['PYTHONPATH'] = ':'.join(filtered_paths)
482        return env_vars
483
484    # pylint: disable=unnecessary-pass
485    # Please keep above disable flag to ensure host_env_check is overriden.
486    def host_env_check(self):
487        """Check that host env has everything we need.
488
489        We actually can assume the host env is fine because we have the same
490        requirements that atest has. Update this to check for android env vars
491        if that changes.
492        """
493        pass
494
495    @staticmethod
496    def _is_missing_exec(executable):
497        """Check if system build executable is available.
498
499        Args:
500            executable: Executable we are checking for.
501
502        Returns:
503            True if executable is missing, False otherwise.
504        """
505        output = shutil.which(executable)
506        if not output:
507            return True
508        # TODO: Check if there is a clever way to determine if system adb is
509        # good enough.
510        root_dir = os.environ.get(constants.ANDROID_BUILD_TOP)
511        return os.path.commonprefix([output, root_dir]) != root_dir
512
513    def get_test_runner_build_reqs(self):
514        """Return the build requirements.
515
516        Returns:
517            Set of build targets.
518        """
519        build_req = self._BUILD_REQ
520        # Use different base build requirements if google-tf is around.
521        if self.module_info.is_module(constants.GTF_MODULE):
522            build_req = {constants.GTF_TARGET}
523        # Always add ATest's own TF target.
524        build_req.add(constants.ATEST_TF_MODULE)
525        # Add adb if we can't find it.
526        for executable in EXEC_DEPENDENCIES:
527            if self._is_missing_exec(executable):
528                build_req.add(executable)
529        return build_req
530
531    # pylint: disable=too-many-branches
532    # pylint: disable=too-many-statements
533    def _parse_extra_args(self, test_infos, extra_args):
534        """Convert the extra args into something tf can understand.
535
536        Args:
537            extra_args: Dict of args
538
539        Returns:
540            Tuple of args to append and args not supported.
541        """
542        args_to_append = []
543        args_not_supported = []
544        for arg in extra_args:
545            if constants.WAIT_FOR_DEBUGGER == arg:
546                args_to_append.append('--wait-for-debugger')
547                continue
548            if constants.DISABLE_INSTALL == arg:
549                args_to_append.append('--disable-target-preparers')
550                continue
551            if constants.SERIAL == arg:
552                args_to_append.append('--serial')
553                args_to_append.append(extra_args[arg])
554                continue
555            if constants.SHARDING == arg:
556                args_to_append.append('--shard-count')
557                args_to_append.append(str(extra_args[arg]))
558                continue
559            if constants.DISABLE_TEARDOWN == arg:
560                args_to_append.append('--disable-teardown')
561                continue
562            if constants.HOST == arg:
563                args_to_append.append('-n')
564                args_to_append.append('--prioritize-host-config')
565                args_to_append.append('--skip-host-arch-check')
566                continue
567            if constants.CUSTOM_ARGS == arg:
568                # We might need to sanitize it prior to appending but for now
569                # let's just treat it like a simple arg to pass on through.
570                args_to_append.extend(extra_args[arg])
571                continue
572            if constants.ALL_ABI == arg:
573                args_to_append.append('--all-abi')
574                continue
575            if constants.DRY_RUN == arg:
576                continue
577            if constants.FLAKES_INFO == arg:
578                continue
579            if constants.INSTANT == arg:
580                args_to_append.append('--enable-parameterized-modules')
581                args_to_append.append('--module-parameter')
582                args_to_append.append('instant_app')
583                continue
584            if constants.USER_TYPE == arg:
585                args_to_append.append('--enable-parameterized-modules')
586                args_to_append.append('--enable-optional-parameterization')
587                args_to_append.append('--module-parameter')
588                args_to_append.append(extra_args[arg])
589                continue
590            if constants.ITERATIONS == arg:
591                args_to_append.append('--retry-strategy')
592                args_to_append.append(constants.ITERATIONS)
593                args_to_append.append('--max-testcase-run-count')
594                args_to_append.append(str(extra_args[arg]))
595                continue
596            if constants.RERUN_UNTIL_FAILURE == arg:
597                args_to_append.append('--retry-strategy')
598                args_to_append.append(constants.RERUN_UNTIL_FAILURE)
599                args_to_append.append('--max-testcase-run-count')
600                args_to_append.append(str(extra_args[arg]))
601                continue
602            if constants.RETRY_ANY_FAILURE == arg:
603                args_to_append.append('--retry-strategy')
604                args_to_append.append(constants.RETRY_ANY_FAILURE)
605                args_to_append.append('--max-testcase-run-count')
606                args_to_append.append(str(extra_args[arg]))
607                continue
608            if constants.COLLECT_TESTS_ONLY == arg:
609                args_to_append.append('--collect-tests-only')
610                continue
611            if constants.TF_DEBUG == arg:
612                print("Please attach process to your IDE...")
613                continue
614            if arg in (constants.TF_TEMPLATE,
615                       constants.TF_EARLY_DEVICE_RELEASE,
616                       constants.INVOCATION_ID,
617                       constants.WORKUNIT_ID):
618                continue
619            args_not_supported.append(arg)
620        # Set exclude instant app annotation for non-instant mode run.
621        if (constants.INSTANT not in extra_args and
622            self._has_instant_app_config(test_infos, self.module_info)):
623            args_to_append.append(constants.TF_TEST_ARG)
624            args_to_append.append(
625                '{tf_class}:{option_name}:{option_value}'.format(
626                    tf_class=constants.TF_AND_JUNIT_CLASS,
627                    option_name=constants.TF_EXCLUDE_ANNOTATE,
628                    option_value=constants.INSTANT_MODE_ANNOTATE))
629        # If test config has config with auto enable parameter, force exclude
630        # those default parameters(ex: instant_app, secondary_user)
631        if '--enable-parameterized-modules' not in args_to_append:
632            for tinfo in test_infos:
633                if self._is_parameter_auto_enabled_cfg(tinfo, self.module_info):
634                    args_to_append.append('--enable-parameterized-modules')
635                    for exclude_parameter in constants.DEFAULT_EXCLUDE_PARAS:
636                        args_to_append.append('--exclude-module-parameters')
637                        args_to_append.append(exclude_parameter)
638                    break
639        return args_to_append, args_not_supported
640
641    def _generate_metrics_folder(self, extra_args):
642        """Generate metrics folder."""
643        metrics_folder = ''
644        if extra_args.get(constants.PRE_PATCH_ITERATIONS):
645            metrics_folder = os.path.join(self.results_dir, 'baseline-metrics')
646        elif extra_args.get(constants.POST_PATCH_ITERATIONS):
647            metrics_folder = os.path.join(self.results_dir, 'new-metrics')
648        return metrics_folder
649
650    def _generate_iterations(self, extra_args):
651        """Generate iterations."""
652        iterations = 1
653        if extra_args.get(constants.PRE_PATCH_ITERATIONS):
654            iterations = extra_args.pop(constants.PRE_PATCH_ITERATIONS)
655        elif extra_args.get(constants.POST_PATCH_ITERATIONS):
656            iterations = extra_args.pop(constants.POST_PATCH_ITERATIONS)
657        return iterations
658
659    def generate_run_commands(self, test_infos, extra_args, port=None):
660        """Generate a single run command from TestInfos.
661
662        Args:
663            test_infos: A set of TestInfo instances.
664            extra_args: A Dict of extra args to append.
665            port: Optional. An int of the port number to send events to. If
666                  None, then subprocess reporter in TF won't try to connect.
667
668        Returns:
669            A list that contains the string of atest tradefed run command.
670            Only one command is returned.
671        """
672        args = self._create_test_args(test_infos)
673        metrics_folder = self._generate_metrics_folder(extra_args)
674
675        # Create a copy of args as more args could be added to the list.
676        test_args = list(args)
677        if port:
678            test_args.extend(['--subprocess-report-port', str(port)])
679        if metrics_folder:
680            test_args.extend(['--metrics-folder', metrics_folder])
681            logging.info('Saved metrics in: %s', metrics_folder)
682        if extra_args.get(constants.INVOCATION_ID, None):
683            test_args.append('--invocation-data invocation_id=%s'
684                             % extra_args[constants.INVOCATION_ID])
685        if extra_args.get(constants.WORKUNIT_ID, None):
686            test_args.append('--invocation-data work_unit_id=%s'
687                             % extra_args[constants.WORKUNIT_ID])
688        # For detailed logs, set TF options log-level/log-level-display as
689        # 'VERBOSE' by default.
690        log_level = 'VERBOSE'
691        test_args.extend(['--log-level-display', log_level])
692        test_args.extend(['--log-level', log_level])
693        # Set no-early-device-release by default to speed up TF teardown time.
694        if not constants.TF_EARLY_DEVICE_RELEASE in extra_args:
695            test_args.extend(['--no-early-device-release'])
696
697        args_to_add, args_not_supported = self._parse_extra_args(test_infos, extra_args)
698
699        # TODO(b/122889707) Remove this after finding the root cause.
700        env_serial = os.environ.get(constants.ANDROID_SERIAL)
701        # Use the env variable ANDROID_SERIAL if it's set by user but only when
702        # the target tests are not deviceless tests.
703        if env_serial and '--serial' not in args_to_add and '-n' not in args_to_add:
704            args_to_add.append("--serial")
705            args_to_add.append(env_serial)
706
707        test_args.extend(args_to_add)
708        if args_not_supported:
709            logging.info('%s does not support the following args %s',
710                         self.EXECUTABLE, args_not_supported)
711
712        # Only need to check one TestInfo to determine if the tests are
713        # configured in TEST_MAPPING.
714        for_test_mapping = test_infos and test_infos[0].from_test_mapping
715        test_args.extend(atest_utils.get_result_server_args(for_test_mapping))
716        self.run_cmd_dict['args'] = ' '.join(test_args)
717        self.run_cmd_dict['tf_customize_template'] = (
718            self._extract_customize_tf_templates(extra_args))
719        return [self._RUN_CMD.format(**self.run_cmd_dict)]
720
721    def _flatten_test_infos(self, test_infos):
722        """Sort and group test_infos by module_name and sort and group filters
723        by class name.
724
725            Example of three test_infos in a set:
726                Module1, {(classA, {})}
727                Module1, {(classB, {Method1})}
728                Module1, {(classB, {Method2}}
729            Becomes a set with one element:
730                Module1, {(ClassA, {}), (ClassB, {Method1, Method2})}
731            Where:
732                  Each line is a test_info namedtuple
733                  {} = Frozenset
734                  () = TestFilter namedtuple
735
736        Args:
737            test_infos: A set of TestInfo namedtuples.
738
739        Returns:
740            A set of TestInfos flattened.
741        """
742        results = set()
743        key = lambda x: x.test_name
744        for module, group in atest_utils.sort_and_group(test_infos, key):
745            # module is a string, group is a generator of grouped TestInfos.
746            # Module Test, so flatten test_infos:
747            no_filters = False
748            filters = set()
749            test_runner = None
750            test_finder = None
751            build_targets = set()
752            data = {}
753            module_args = []
754            for test_info_i in group:
755                data.update(test_info_i.data)
756                # Extend data with constants.TI_MODULE_ARG instead of
757                # overwriting.
758                module_args.extend(test_info_i.data.get(
759                    constants.TI_MODULE_ARG, []))
760                test_runner = test_info_i.test_runner
761                test_finder = test_info_i.test_finder
762                build_targets |= test_info_i.build_targets
763                test_filters = test_info_i.data.get(constants.TI_FILTER)
764                if not test_filters or no_filters:
765                    # test_info wants whole module run, so hardcode no filters.
766                    no_filters = True
767                    filters = set()
768                    continue
769                filters |= test_filters
770            if module_args:
771                data[constants.TI_MODULE_ARG] = module_args
772            data[constants.TI_FILTER] = self._flatten_test_filters(filters)
773            results.add(
774                test_info.TestInfo(test_name=module,
775                                   test_runner=test_runner,
776                                   test_finder=test_finder,
777                                   build_targets=build_targets,
778                                   data=data))
779        return results
780
781    @staticmethod
782    def _flatten_test_filters(filters):
783        """Sort and group test_filters by class_name.
784
785            Example of three test_filters in a frozenset:
786                classA, {}
787                classB, {Method1}
788                classB, {Method2}
789            Becomes a frozenset with these elements:
790                classA, {}
791                classB, {Method1, Method2}
792            Where:
793                Each line is a TestFilter namedtuple
794                {} = Frozenset
795
796        Args:
797            filters: A frozenset of test_filters.
798
799        Returns:
800            A frozenset of test_filters flattened.
801        """
802        results = set()
803        key = lambda x: x.class_name
804        for class_name, group in atest_utils.sort_and_group(filters, key):
805            # class_name is a string, group is a generator of TestFilters
806            assert class_name is not None
807            methods = set()
808            for test_filter in group:
809                if not test_filter.methods:
810                    # Whole class should be run
811                    methods = set()
812                    break
813                methods |= test_filter.methods
814            results.add(test_info.TestFilter(class_name, frozenset(methods)))
815        return frozenset(results)
816
817    def _create_test_args(self, test_infos):
818        """Compile TF command line args based on the given test infos.
819
820        Args:
821            test_infos: A set of TestInfo instances.
822
823        Returns: A list of TF arguments to run the tests.
824        """
825        args = []
826        if not test_infos:
827            return []
828
829        test_infos = self._flatten_test_infos(test_infos)
830        has_integration_test = False
831        for info in test_infos:
832            # Integration test exists in TF's jar, so it must have the option
833            # if it's integration finder.
834            if info.test_finder in _INTEGRATION_FINDERS:
835                has_integration_test = True
836            # For non-paramertize test module, use --include-filter, but for
837            # tests which have auto enable paramertize config use --module
838            # instead.
839            if self._is_parameter_auto_enabled_cfg(info, self.module_info):
840                args.extend([constants.TF_MODULE_FILTER, info.test_name])
841            else:
842                args.extend([constants.TF_INCLUDE_FILTER, info.test_name])
843            filters = set()
844            for test_filter in info.data.get(constants.TI_FILTER, []):
845                filters.update(test_filter.to_set_of_tf_strings())
846            for test_filter in filters:
847                filter_arg = constants.TF_ATEST_INCLUDE_FILTER_VALUE_FMT.format(
848                    test_name=info.test_name, test_filter=test_filter)
849                args.extend([constants.TF_ATEST_INCLUDE_FILTER, filter_arg])
850            for option in info.data.get(constants.TI_MODULE_ARG, []):
851                if constants.TF_INCLUDE_FILTER_OPTION == option[0]:
852                    suite_filter = (
853                        constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
854                            test_name=info.test_name, option_value=option[1]))
855                    args.extend([constants.TF_INCLUDE_FILTER, suite_filter])
856                elif constants.TF_EXCLUDE_FILTER_OPTION == option[0]:
857                    suite_filter = (
858                        constants.TF_SUITE_FILTER_ARG_VALUE_FMT.format(
859                            test_name=info.test_name, option_value=option[1]))
860                    args.extend([constants.TF_EXCLUDE_FILTER, suite_filter])
861                else:
862                    module_arg = (
863                        constants.TF_MODULE_ARG_VALUE_FMT.format(
864                            test_name=info.test_name, option_name=option[0],
865                            option_value=option[1]))
866                    args.extend([constants.TF_MODULE_ARG, module_arg])
867        # TODO (b/141090547) Pass the config path to TF to load configs.
868        # Compile option in TF if finder is not INTEGRATION or not set.
869        if not has_integration_test:
870            args.append(constants.TF_SKIP_LOADING_CONFIG_JAR)
871        return args
872
873    def _extract_rerun_options(self, extra_args):
874        """Extract rerun options to a string for output.
875
876        Args:
877            extra_args: Dict of extra args for test runners to use.
878
879        Returns: A string of rerun options.
880        """
881        extracted_options = ['{} {}'.format(arg, extra_args[arg])
882                             for arg in extra_args
883                             if arg in self._RERUN_OPTION_GROUP]
884        return ' '.join(extracted_options)
885
886    def _extract_customize_tf_templates(self, extra_args):
887        """Extract tradefed template options to a string for output.
888
889        Args:
890            extra_args: Dict of extra args for test runners to use.
891
892        Returns: A string of tradefed template options.
893        """
894        return ' '.join(['--template:map %s'
895                         % x for x in extra_args.get(constants.TF_TEMPLATE, [])])
896
897    def _handle_log_associations(self, event_handlers):
898        """Handle TF's log associations information data.
899
900        log_association dict:
901        {'loggedFile': '/tmp/serial-util11375755456514097276.ser',
902         'dataName': 'device_logcat_setup_127.0.0.1:58331',
903         'time': 1602038599.856113},
904
905        Args:
906            event_handlers: Dict of {socket_object:EventHandler}.
907
908        """
909        log_associations = []
910        for _, event_handler in event_handlers.items():
911            if event_handler.log_associations:
912                log_associations += event_handler.log_associations
913        device_test_end_log_time = ''
914        device_teardown_log_time = ''
915        for log_association in log_associations:
916            if 'device_logcat_test' in log_association.get('dataName', ''):
917                device_test_end_log_time = log_association.get('time')
918            if 'device_logcat_teardown' in log_association.get('dataName', ''):
919                device_teardown_log_time = log_association.get('time')
920        if device_test_end_log_time and device_teardown_log_time:
921            teardowntime = (float(device_teardown_log_time) -
922                            float(device_test_end_log_time))
923            logging.debug('TF logcat teardown time=%s seconds.', teardowntime)
924            metrics.LocalDetectEvent(
925                detect_type=constants.DETECT_TYPE_TF_TEARDOWN_LOGCAT,
926                result=int(teardowntime))
927
928    @staticmethod
929    def _has_instant_app_config(test_infos, mod_info):
930        """Check if one of the input tests defined instant app mode in config.
931
932        Args:
933            test_infos: A set of TestInfo instances.
934            mod_info: ModuleInfo object.
935
936        Returns: True if one of the tests set up instant app mode.
937        """
938        for tinfo in test_infos:
939            test_config, _ = test_finder_utils.get_test_config_and_srcs(
940                tinfo, mod_info)
941            if test_config:
942                parameters = atest_utils.get_config_parameter(test_config)
943                if constants.TF_PARA_INSTANT_APP in parameters:
944                    return True
945        return False
946
947    @staticmethod
948    def _is_parameter_auto_enabled_cfg(tinfo, mod_info):
949        """Check if input tests contains auto enable support parameters.
950
951        Args:
952            test_infos: A set of TestInfo instances.
953            mod_info: ModuleInfo object.
954
955        Returns: True if input test has parameter setting which is not in the
956                 exclude list.
957        """
958        test_config, _ = test_finder_utils.get_test_config_and_srcs(
959            tinfo, mod_info)
960        if test_config:
961            parameters = atest_utils.get_config_parameter(test_config)
962            if (parameters - constants.DEFAULT_EXCLUDE_PARAS
963                - constants.DEFAULT_EXCLUDE_NOT_PARAS):
964                return True
965        return False
966