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"""
16Utility functions for atest.
17"""
18
19
20# pylint: disable=import-outside-toplevel
21
22from __future__ import print_function
23
24import hashlib
25import itertools
26import json
27import logging
28import os
29import pickle
30import re
31import shutil
32import subprocess
33import sys
34
35import atest_decorator
36import atest_error
37import constants
38
39# b/147562331 only occurs when running atest in source code. We don't encourge
40# the users to manually "pip3 install protobuf", therefore when the exception
41# occurs, we don't collect data and the tab completion is for args is silence.
42try:
43    from metrics import metrics_base
44    from metrics import metrics_utils
45except ModuleNotFoundError:
46    # This exception occurs only when invoking atest in source code.
47    print("You shouldn't see this message unless you ran 'atest-src'."
48          "To resolve the issue, please run:\n\t{}\n"
49          "and try again.".format('pip3 install protobuf'))
50    sys.exit(constants.IMPORT_FAILURE)
51
52_BASH_RESET_CODE = '\033[0m\n'
53# Arbitrary number to limit stdout for failed runs in _run_limited_output.
54# Reason for its use is that the make command itself has its own carriage
55# return output mechanism that when collected line by line causes the streaming
56# full_output list to be extremely large.
57_FAILED_OUTPUT_LINE_LIMIT = 100
58# Regular expression to match the start of a ninja compile:
59# ex: [ 99% 39710/39711]
60_BUILD_COMPILE_STATUS = re.compile(r'\[\s*(\d{1,3}%\s+)?\d+/\d+\]')
61_BUILD_FAILURE = 'FAILED: '
62CMD_RESULT_PATH = os.path.join(os.environ.get(constants.ANDROID_BUILD_TOP,
63                                              os.getcwd()),
64                               'tools/tradefederation/core/atest/test_data',
65                               'test_commands.json')
66BUILD_TOP_HASH = hashlib.md5(os.environ.get(constants.ANDROID_BUILD_TOP, '').
67                             encode()).hexdigest()
68TEST_INFO_CACHE_ROOT = os.path.join(os.path.expanduser('~'), '.atest',
69                                    'info_cache', BUILD_TOP_HASH[:8])
70_DEFAULT_TERMINAL_WIDTH = 80
71_DEFAULT_TERMINAL_HEIGHT = 25
72_BUILD_CMD = 'build/soong/soong_ui.bash'
73_FIND_MODIFIED_FILES_CMDS = (
74    "cd {};"
75    "local_branch=$(git rev-parse --abbrev-ref HEAD);"
76    "remote_branch=$(git branch -r | grep '\\->' | awk '{{print $1}}');"
77    # Get the number of commits from local branch to remote branch.
78    "ahead=$(git rev-list --left-right --count $local_branch...$remote_branch "
79    "| awk '{{print $1}}');"
80    # Get the list of modified files from HEAD to previous $ahead generation.
81    "git diff HEAD~$ahead --name-only")
82
83
84def get_build_cmd():
85    """Compose build command with no-absolute path and flag "--make-mode".
86
87    Returns:
88        A list of soong build command.
89    """
90    make_cmd = ('%s/%s' %
91                (os.path.relpath(os.environ.get(
92                    constants.ANDROID_BUILD_TOP, os.getcwd()), os.getcwd()),
93                 _BUILD_CMD))
94    return [make_cmd, '--make-mode']
95
96
97def _capture_fail_section(full_log):
98    """Return the error message from the build output.
99
100    Args:
101        full_log: List of strings representing full output of build.
102
103    Returns:
104        capture_output: List of strings that are build errors.
105    """
106    am_capturing = False
107    capture_output = []
108    for line in full_log:
109        if am_capturing and _BUILD_COMPILE_STATUS.match(line):
110            break
111        if am_capturing or line.startswith(_BUILD_FAILURE):
112            capture_output.append(line)
113            am_capturing = True
114            continue
115    return capture_output
116
117
118def _run_limited_output(cmd, env_vars=None):
119    """Runs a given command and streams the output on a single line in stdout.
120
121    Args:
122        cmd: A list of strings representing the command to run.
123        env_vars: Optional arg. Dict of env vars to set during build.
124
125    Raises:
126        subprocess.CalledProcessError: When the command exits with a non-0
127            exitcode.
128    """
129    # Send stderr to stdout so we only have to deal with a single pipe.
130    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
131                            stderr=subprocess.STDOUT, env=env_vars)
132    sys.stdout.write('\n')
133    term_width, _ = get_terminal_size()
134    white_space = " " * int(term_width)
135    full_output = []
136    while proc.poll() is None:
137        line = proc.stdout.readline().decode('utf-8')
138        # Readline will often return empty strings.
139        if not line:
140            continue
141        full_output.append(line)
142        # Trim the line to the width of the terminal.
143        # Note: Does not handle terminal resizing, which is probably not worth
144        #       checking the width every loop.
145        if len(line) >= term_width:
146            line = line[:term_width - 1]
147        # Clear the last line we outputted.
148        sys.stdout.write('\r%s\r' % white_space)
149        sys.stdout.write('%s' % line.strip())
150        sys.stdout.flush()
151    # Reset stdout (on bash) to remove any custom formatting and newline.
152    sys.stdout.write(_BASH_RESET_CODE)
153    sys.stdout.flush()
154    # Wait for the Popen to finish completely before checking the returncode.
155    proc.wait()
156    if proc.returncode != 0:
157        # Parse out the build error to output.
158        output = _capture_fail_section(full_output)
159        if not output:
160            output = full_output
161        if len(output) >= _FAILED_OUTPUT_LINE_LIMIT:
162            output = output[-_FAILED_OUTPUT_LINE_LIMIT:]
163        output = 'Output (may be trimmed):\n%s' % ''.join(output)
164        raise subprocess.CalledProcessError(proc.returncode, cmd, output)
165
166
167def build(build_targets, verbose=False, env_vars=None):
168    """Shell out and make build_targets.
169
170    Args:
171        build_targets: A set of strings of build targets to make.
172        verbose: Optional arg. If True output is streamed to the console.
173                 If False, only the last line of the build output is outputted.
174        env_vars: Optional arg. Dict of env vars to set during build.
175
176    Returns:
177        Boolean of whether build command was successful, True if nothing to
178        build.
179    """
180    if not build_targets:
181        logging.debug('No build targets, skipping build.')
182        return True
183    full_env_vars = os.environ.copy()
184    if env_vars:
185        full_env_vars.update(env_vars)
186    print('\n%s\n%s' % (colorize("Building Dependencies...", constants.CYAN),
187                        ', '.join(build_targets)))
188    logging.debug('Building Dependencies: %s', ' '.join(build_targets))
189    cmd = get_build_cmd() + list(build_targets)
190    logging.debug('Executing command: %s', cmd)
191    try:
192        if verbose:
193            subprocess.check_call(cmd, stderr=subprocess.STDOUT,
194                                  env=full_env_vars)
195        else:
196            # TODO: Save output to a log file.
197            _run_limited_output(cmd, env_vars=full_env_vars)
198        logging.info('Build successful')
199        return True
200    except subprocess.CalledProcessError as err:
201        logging.error('Error building: %s', build_targets)
202        if err.output:
203            logging.error(err.output)
204        return False
205
206
207def _can_upload_to_result_server():
208    """Return True if we can talk to result server."""
209    # TODO: Also check if we have a slow connection to result server.
210    if constants.RESULT_SERVER:
211        try:
212            from urllib.request import urlopen
213            urlopen(constants.RESULT_SERVER,
214                    timeout=constants.RESULT_SERVER_TIMEOUT).close()
215            return True
216        # pylint: disable=broad-except
217        except Exception as err:
218            logging.debug('Talking to result server raised exception: %s', err)
219    return False
220
221
222def get_result_server_args(for_test_mapping=False):
223    """Return list of args for communication with result server.
224
225    Args:
226        for_test_mapping: True if the test run is for Test Mapping to include
227            additional reporting args. Default is False.
228    """
229    # TODO (b/147644460) Temporarily disable Sponge V1 since it will be turned
230    # down.
231    if _can_upload_to_result_server():
232        if for_test_mapping:
233            return (constants.RESULT_SERVER_ARGS +
234                    constants.TEST_MAPPING_RESULT_SERVER_ARGS)
235        return constants.RESULT_SERVER_ARGS
236    return []
237
238
239def sort_and_group(iterable, key):
240    """Sort and group helper function."""
241    return itertools.groupby(sorted(iterable, key=key), key=key)
242
243
244def is_test_mapping(args):
245    """Check if the atest command intends to run tests in test mapping.
246
247    When atest runs tests in test mapping, it must have at most one test
248    specified. If a test is specified, it must be started with  `:`,
249    which means the test value is a test group name in TEST_MAPPING file, e.g.,
250    `:postsubmit`.
251
252    If any test mapping options is specified, the atest command must also be
253    set to run tests in test mapping files.
254
255    Args:
256        args: arg parsed object.
257
258    Returns:
259        True if the args indicates atest shall run tests in test mapping. False
260        otherwise.
261    """
262    return (
263        args.test_mapping or
264        args.include_subdirs or
265        not args.tests or
266        (len(args.tests) == 1 and args.tests[0][0] == ':'))
267
268@atest_decorator.static_var("cached_has_colors", {})
269def _has_colors(stream):
270    """Check the output stream is colorful.
271
272    Args:
273        stream: The standard file stream.
274
275    Returns:
276        True if the file stream can interpreter the ANSI color code.
277    """
278    cached_has_colors = _has_colors.cached_has_colors
279    if stream in cached_has_colors:
280        return cached_has_colors[stream]
281    cached_has_colors[stream] = True
282    # Following from Python cookbook, #475186
283    if not hasattr(stream, "isatty"):
284        cached_has_colors[stream] = False
285        return False
286    if not stream.isatty():
287        # Auto color only on TTYs
288        cached_has_colors[stream] = False
289        return False
290    try:
291        import curses
292        curses.setupterm()
293        cached_has_colors[stream] = curses.tigetnum("colors") > 2
294    # pylint: disable=broad-except
295    except Exception as err:
296        logging.debug('Checking colorful raised exception: %s', err)
297        cached_has_colors[stream] = False
298    return cached_has_colors[stream]
299
300
301def colorize(text, color, highlight=False):
302    """ Convert to colorful string with ANSI escape code.
303
304    Args:
305        text: A string to print.
306        color: ANSI code shift for colorful print. They are defined
307               in constants_default.py.
308        highlight: True to print with highlight.
309
310    Returns:
311        Colorful string with ANSI escape code.
312    """
313    clr_pref = '\033[1;'
314    clr_suff = '\033[0m'
315    has_colors = _has_colors(sys.stdout)
316    if has_colors:
317        if highlight:
318            ansi_shift = 40 + color
319        else:
320            ansi_shift = 30 + color
321        clr_str = "%s%dm%s%s" % (clr_pref, ansi_shift, text, clr_suff)
322    else:
323        clr_str = text
324    return clr_str
325
326
327def colorful_print(text, color, highlight=False, auto_wrap=True):
328    """Print out the text with color.
329
330    Args:
331        text: A string to print.
332        color: ANSI code shift for colorful print. They are defined
333               in constants_default.py.
334        highlight: True to print with highlight.
335        auto_wrap: If True, Text wraps while print.
336    """
337    output = colorize(text, color, highlight)
338    if auto_wrap:
339        print(output)
340    else:
341        print(output, end="")
342
343
344def get_terminal_size():
345    """Get terminal size and return a tuple.
346
347    Returns:
348        2 integers: the size of X(columns) and Y(lines/rows).
349    """
350    # Determine the width of the terminal. We'll need to clear this many
351    # characters when carriage returning. Set default value as 80.
352    columns, rows = shutil.get_terminal_size(
353        fallback=(_DEFAULT_TERMINAL_WIDTH,
354                  _DEFAULT_TERMINAL_HEIGHT))
355    return columns, rows
356
357
358def is_external_run():
359    # TODO(b/133905312): remove this function after aidegen calling
360    #       metrics_base.get_user_type directly.
361    """Check is external run or not.
362
363    Determine the internal user by passing at least one check:
364      - whose git mail domain is from google
365      - whose hostname is from google
366    Otherwise is external user.
367
368    Returns:
369        True if this is an external run, False otherwise.
370    """
371    return metrics_base.get_user_type() == metrics_base.EXTERNAL_USER
372
373
374def print_data_collection_notice():
375    """Print the data collection notice."""
376    anonymous = ''
377    user_type = 'INTERNAL'
378    if metrics_base.get_user_type() == metrics_base.EXTERNAL_USER:
379        anonymous = ' anonymous'
380        user_type = 'EXTERNAL'
381    notice = ('  We collect%s usage statistics in accordance with our Content '
382              'Licenses (%s), Contributor License Agreement (%s), Privacy '
383              'Policy (%s) and Terms of Service (%s).'
384             ) % (anonymous,
385                  constants.CONTENT_LICENSES_URL,
386                  constants.CONTRIBUTOR_AGREEMENT_URL[user_type],
387                  constants.PRIVACY_POLICY_URL,
388                  constants.TERMS_SERVICE_URL
389                 )
390    print(delimiter('=', 18, prenl=1))
391    colorful_print("Notice:", constants.RED)
392    colorful_print("%s" % notice, constants.GREEN)
393    print(delimiter('=', 18, postnl=1))
394
395
396def handle_test_runner_cmd(input_test, test_cmds, do_verification=False,
397                           result_path=CMD_RESULT_PATH):
398    """Handle the runner command of input tests.
399
400    Args:
401        input_test: A string of input tests pass to atest.
402        test_cmds: A list of strings for running input tests.
403        do_verification: A boolean to indicate the action of this method.
404                         True: Do verification without updating result map and
405                               raise DryRunVerificationError if verifying fails.
406                         False: Update result map, if the former command is
407                                different with current command, it will confirm
408                                with user if they want to update or not.
409        result_path: The file path for saving result.
410    """
411    full_result_content = {}
412    if os.path.isfile(result_path):
413        with open(result_path) as json_file:
414            full_result_content = json.load(json_file)
415    former_test_cmds = full_result_content.get(input_test, [])
416    if not _are_identical_cmds(test_cmds, former_test_cmds):
417        if do_verification:
418            raise atest_error.DryRunVerificationError(
419                'Dry run verification failed, former commands: {}'.format(
420                    former_test_cmds))
421        if former_test_cmds:
422            # If former_test_cmds is different from test_cmds, ask users if they
423            # are willing to update the result.
424            print('Former cmds = %s' % former_test_cmds)
425            print('Current cmds = %s' % test_cmds)
426            try:
427                from distutils import util
428                if not util.strtobool(
429                        input('Do you want to update former result '
430                              'with the latest one?(Y/n)')):
431                    print('SKIP updating result!!!')
432                    return
433            except ValueError:
434                # Default action is updating the command result of the
435                # input_test. If the user input is unrecognizable telling yes
436                # or no, "Y" is implicitly applied.
437                pass
438    else:
439        # If current commands are the same as the formers, no need to update
440        # result.
441        return
442    full_result_content[input_test] = test_cmds
443    with open(result_path, 'w') as outfile:
444        json.dump(full_result_content, outfile, indent=0)
445        print('Save result mapping to %s' % result_path)
446
447
448def _are_identical_cmds(current_cmds, former_cmds):
449    """Tell two commands are identical. Note that '--atest-log-file-path' is not
450    considered a critical argument, therefore, it will be removed during
451    the comparison. Also, atest can be ran in any place, so verifying relative
452    path is regardless as well.
453
454    Args:
455        current_cmds: A list of strings for running input tests.
456        former_cmds: A list of strings recorded from the previous run.
457
458    Returns:
459        True if both commands are identical, False otherwise.
460    """
461    def _normalize(cmd_list):
462        """Method that normalize commands.
463
464        Args:
465            cmd_list: A list with one element. E.g. ['cmd arg1 arg2 True']
466
467        Returns:
468            A list with elements. E.g. ['cmd', 'arg1', 'arg2', 'True']
469        """
470        _cmd = ''.join(cmd_list).split()
471        for cmd in _cmd:
472            if cmd.startswith('--atest-log-file-path'):
473                _cmd.remove(cmd)
474                continue
475            if _BUILD_CMD in cmd:
476                _cmd.remove(cmd)
477                _cmd.append(os.path.join('./', _BUILD_CMD))
478                continue
479        return _cmd
480
481    _current_cmds = _normalize(current_cmds)
482    _former_cmds = _normalize(former_cmds)
483    # Always sort cmd list to make it comparable.
484    _current_cmds.sort()
485    _former_cmds.sort()
486    return _current_cmds == _former_cmds
487
488def _get_hashed_file_name(main_file_name):
489    """Convert the input string to a md5-hashed string. If file_extension is
490       given, returns $(hashed_string).$(file_extension), otherwise
491       $(hashed_string).cache.
492
493    Args:
494        main_file_name: The input string need to be hashed.
495
496    Returns:
497        A string as hashed file name with .cache file extension.
498    """
499    hashed_fn = hashlib.md5(str(main_file_name).encode())
500    hashed_name = hashed_fn.hexdigest()
501    return hashed_name + '.cache'
502
503def get_test_info_cache_path(test_reference, cache_root=TEST_INFO_CACHE_ROOT):
504    """Get the cache path of the desired test_infos.
505
506    Args:
507        test_reference: A string of the test.
508        cache_root: Folder path where stores caches.
509
510    Returns:
511        A string of the path of test_info cache.
512    """
513    return os.path.join(cache_root,
514                        _get_hashed_file_name(test_reference))
515
516def update_test_info_cache(test_reference, test_infos,
517                           cache_root=TEST_INFO_CACHE_ROOT):
518    """Update cache content which stores a set of test_info objects through
519       pickle module, each test_reference will be saved as a cache file.
520
521    Args:
522        test_reference: A string referencing a test.
523        test_infos: A set of TestInfos.
524        cache_root: Folder path for saving caches.
525    """
526    if not os.path.isdir(cache_root):
527        os.makedirs(cache_root)
528    cache_path = get_test_info_cache_path(test_reference, cache_root)
529    # Save test_info to files.
530    try:
531        with open(cache_path, 'wb') as test_info_cache_file:
532            logging.debug('Saving cache %s.', cache_path)
533            pickle.dump(test_infos, test_info_cache_file, protocol=2)
534    except (pickle.PicklingError, TypeError, IOError) as err:
535        # Won't break anything, just log this error, and collect the exception
536        # by metrics.
537        logging.debug('Exception raised: %s', err)
538        metrics_utils.handle_exc_and_send_exit_event(
539            constants.ACCESS_CACHE_FAILURE)
540
541
542def load_test_info_cache(test_reference, cache_root=TEST_INFO_CACHE_ROOT):
543    """Load cache by test_reference to a set of test_infos object.
544
545    Args:
546        test_reference: A string referencing a test.
547        cache_root: Folder path for finding caches.
548
549    Returns:
550        A list of TestInfo namedtuple if cache found, else None.
551    """
552    cache_file = get_test_info_cache_path(test_reference, cache_root)
553    if os.path.isfile(cache_file):
554        logging.debug('Loading cache %s.', cache_file)
555        try:
556            with open(cache_file, 'rb') as config_dictionary_file:
557                return pickle.load(config_dictionary_file, encoding='utf-8')
558        except (pickle.UnpicklingError,
559                ValueError,
560                TypeError,
561                EOFError,
562                IOError) as err:
563            # Won't break anything, just remove the old cache, log this error,
564            # and collect the exception by metrics.
565            logging.debug('Exception raised: %s', err)
566            os.remove(cache_file)
567            metrics_utils.handle_exc_and_send_exit_event(
568                constants.ACCESS_CACHE_FAILURE)
569    return None
570
571def clean_test_info_caches(tests, cache_root=TEST_INFO_CACHE_ROOT):
572    """Clean caches of input tests.
573
574    Args:
575        tests: A list of test references.
576        cache_root: Folder path for finding caches.
577    """
578    for test in tests:
579        cache_file = get_test_info_cache_path(test, cache_root)
580        if os.path.isfile(cache_file):
581            logging.debug('Removing cache: %s', cache_file)
582            try:
583                os.remove(cache_file)
584            except IOError as err:
585                logging.debug('Exception raised: %s', err)
586                metrics_utils.handle_exc_and_send_exit_event(
587                    constants.ACCESS_CACHE_FAILURE)
588
589def get_modified_files(root_dir):
590    """Get the git modified files. The git path here is git top level of
591    the root_dir. It's inevitable to utilise different commands to fulfill
592    2 scenario:
593        1. locate unstaged/staged files
594        2. locate committed files but not yet merged.
595    the 'git_status_cmd' fulfils the former while the 'find_modified_files'
596    fulfils the latter.
597
598    Args:
599        root_dir: the root where it starts finding.
600
601    Returns:
602        A set of modified files altered since last commit.
603    """
604    modified_files = set()
605    try:
606        find_git_cmd = 'cd {}; git rev-parse --show-toplevel'.format(root_dir)
607        git_paths = subprocess.check_output(
608            find_git_cmd, shell=True).decode().splitlines()
609        for git_path in git_paths:
610            # Find modified files from git working tree status.
611            git_status_cmd = ("repo forall {} -c git status --short | "
612                              "awk '{{print $NF}}'").format(git_path)
613            modified_wo_commit = subprocess.check_output(
614                git_status_cmd, shell=True).decode().rstrip().splitlines()
615            for change in modified_wo_commit:
616                modified_files.add(
617                    os.path.normpath('{}/{}'.format(git_path, change)))
618            # Find modified files that are committed but not yet merged.
619            find_modified_files = _FIND_MODIFIED_FILES_CMDS.format(git_path)
620            commit_modified_files = subprocess.check_output(
621                find_modified_files, shell=True).decode().splitlines()
622            for line in commit_modified_files:
623                modified_files.add(os.path.normpath('{}/{}'.format(
624                    git_path, line)))
625    except (OSError, subprocess.CalledProcessError) as err:
626        logging.debug('Exception raised: %s', err)
627    return modified_files
628
629def delimiter(char, length=_DEFAULT_TERMINAL_WIDTH, prenl=0, postnl=0):
630    """A handy delimiter printer.
631
632    Args:
633        char: A string used for delimiter.
634        length: An integer for the replication.
635        prenl: An integer that insert '\n' before delimiter.
636        postnl: An integer that insert '\n' after delimiter.
637
638    Returns:
639        A string of delimiter.
640    """
641    return prenl * '\n' + char * length + postnl * '\n'
642