1#!/usr/bin/env python3
2
3# Copyright (C) 2023 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Script for running git-based Mobly tests locally.
18
19Example:
20    - Run an Android platform test module.
21    local_mobly_runner.py -m my_test_module
22
23    - Run an Android platform test module. Build the module and install test
24      APKs before running the test.
25    local_mobly_runner.py -m my_test_module -b -i
26
27    - Run an Android platform test module with specific Android devices.
28    local_mobly_runner.py -m my_test_module -s DEV00001,DEV00002
29
30    - Run a list of zipped executable Mobly packages
31    local_mobly_runner.py -p test_pkg1,test_pkg2,test_pkg3
32
33    - Install and run a test binary from a Python wheel
34    local_mobly_runner.py -w my-test-0.1-py3-none-any.whl --bin test_suite_a
35
36Please run `local_mobly_runner.py -h` for a full list of options.
37"""
38
39import argparse
40import json
41import os
42from pathlib import Path
43import platform
44import shutil
45import subprocess
46import sys
47import tempfile
48from typing import List, Optional, Tuple
49import zipfile
50
51_LOCAL_SETUP_INSTRUCTIONS = (
52    '\n\tcd <repo_root>; set -a; source build/envsetup.sh; set +a; lunch'
53    ' <target>'
54)
55_DEFAULT_MOBLY_LOGPATH = Path('/tmp/logs/mobly')
56_DEFAULT_TESTBED = 'LocalTestBed'
57
58_tempdirs = []
59_tempfiles = []
60
61
62def _padded_print(line: str) -> None:
63    print(f'\n-----{line}-----\n')
64
65
66def _parse_args() -> argparse.Namespace:
67    """Parses command line args."""
68    parser = argparse.ArgumentParser(
69        formatter_class=argparse.RawDescriptionHelpFormatter,
70        description=__doc__)
71    group1 = parser.add_mutually_exclusive_group(required=True)
72    group1.add_argument(
73        '-m', '--module',
74        help='The Android platform build module of the test to run.'
75    )
76    group1.add_argument(
77        '-p', '--packages',
78        help=(
79            'A comma-delimited list of test packages to run. The packages '
80            'should be directly executable by the Python interpreter. If '
81            'the package includes a requirements.txt file, deps will '
82            'automatically be installed.'
83        )
84    )
85    group1.add_argument(
86        '-w', '--wheel',
87        help=(
88            'A Python wheel (.whl) containing one or more Mobly test scripts. '
89            'Does not support the --novenv option.'
90        )
91    )
92    group1.add_argument(
93        '-t',
94        '--test_paths',
95        help=(
96            'A comma-delimited list of test paths to run directly. Implies '
97            'the --novenv option.'
98        ),
99    )
100
101    parser.add_argument(
102        '-b',
103        '--build',
104        action='store_true',
105        help='Build/rebuild the specified module. Requires the -m option.',
106    )
107    parser.add_argument(
108        '-i',
109        '--install_apks',
110        action='store_true',
111        help=(
112            'Install all APKs associated with the module to all specified'
113            ' devices. Does not support the -t option.'
114        ),
115    )
116    parser.add_argument(
117        '-s',
118        '--serials',
119        help=(
120            'Specify the devices to test with a comma-delimited list of device '
121            'serials. If --config is also specified, this option will only be '
122            'used to select the devices to install APKs.'
123        ),
124    )
125    parser.add_argument(
126        '-c', '--config', help='Provide a custom Mobly config for the test.'
127    )
128    parser.add_argument('-tb', '--test_bed',
129                        default=_DEFAULT_TESTBED,
130                        help='Select the testbed for the test. If left '
131                             f'unspecified, "{_DEFAULT_TESTBED}" will be '
132                             'selected by default.')
133    parser.add_argument('-lp', '--log_path',
134                        help='Specify a path to store logs.')
135
136    parser.add_argument(
137        '--tests',
138        nargs='+',
139        type=str,
140        metavar='TEST_CLASS[.TEST_CASE]',
141        help=(
142            'A list of test classes and optional tests to execute within the '
143            'package or file. E.g. `--tests TestClassA TestClassB.test_b` '
144            'would run all of test class TestClassA, but only test_b in '
145            'TestClassB. This option cannot be used if multiple packages/test '
146            'paths are specified.'
147        ),
148    )
149    parser.add_argument(
150        '--bin',
151        help=(
152            'Name of the binary to run in the installed wheel. Must be '
153            'specified alongside the --wheel option.'
154        ),
155    )
156    parser.add_argument(
157        '--novenv',
158        action='store_true',
159        help=(
160            "Run directly in the host's system Python, without setting up a "
161            'virtualenv.'
162        ),
163    )
164    args = parser.parse_args()
165    if args.build and not args.module:
166        parser.error('Option --build requires --module to be specified.')
167    if args.wheel:
168        if args.novenv:
169            parser.error('Option --novenv cannot be used with --wheel.')
170        if not args.bin:
171            parser.error('Option --wheel requires --bin to be specified.')
172    if args.bin:
173        if not args.wheel:
174            parser.error('Option --bin requires --wheel to be specified.')
175    if args.install_apks and args.test_paths:
176        parser.error('Option --install_apks cannot be used with --test_paths.')
177    if args.tests is not None:
178        multiple_packages = (args.packages is not None
179                             and len(args.packages.split(',')) > 1)
180        multiple_test_paths = (args.test_paths is not None
181                               and len(args.test_paths.split(',')) > 1)
182        if multiple_packages or multiple_test_paths:
183            parser.error(
184                'Option --tests cannot be used if multiple --packages or '
185                '--test_paths are specified.'
186            )
187
188    args.novenv = args.novenv or (args.test_paths is not None)
189    return args
190
191
192def _build_module(module: str) -> None:
193    """Builds the specified module."""
194    _padded_print(f'Building test module {module}.')
195    try:
196        subprocess.check_call(f'm -j {module}', shell=True,
197                              executable='/bin/bash')
198    except subprocess.CalledProcessError as e:
199        if e.returncode == 127:
200            # `m` command not found
201            print(
202                '`m` command not found. Please set up your local environment '
203                f'with {_LOCAL_SETUP_INSTRUCTIONS}.'
204            )
205        else:
206            print(f'Failed to build module {module}.')
207        exit(1)
208
209
210def _get_module_artifacts(module: str) -> List[str]:
211    """Return the list of artifacts generated from a module."""
212    try:
213        outmod_paths = (
214            subprocess.check_output(
215                f'outmod {module}', shell=True, executable='/bin/bash'
216            )
217            .decode('utf-8')
218            .splitlines()
219        )
220    except subprocess.CalledProcessError as e:
221        if e.returncode == 127:
222            # `outmod` command not found
223            print(
224                '`outmod` command not found. Please set up your local '
225                f'environment with {_LOCAL_SETUP_INSTRUCTIONS}.'
226            )
227        if str(e.output).startswith('Could not find module'):
228            print(
229                f'Cannot find the build output of module {module}. Ensure that '
230                'the module list is up-to-date with `refreshmod`.'
231            )
232        exit(1)
233
234    for path in outmod_paths:
235        if not os.path.isfile(path):
236            print(
237                f'Declared file {path} does not exist. Please build your '
238                'module with the -b option.'
239            )
240            exit(1)
241
242    return outmod_paths
243
244
245def _extract_test_resources(
246        args: argparse.Namespace,
247) -> Tuple[List[str], List[str], List[str]]:
248    """Extract test resources from the given test module or package.
249
250    Args:
251      args: Parsed command-line args.
252
253    Returns:
254      Tuple of (mobly_bins, requirement_files, test_apks).
255    """
256    _padded_print('Resolving test resources.')
257    mobly_bins = []
258    requirements_files = []
259    test_apks = []
260    if args.test_paths:
261        mobly_bins.extend(args.test_paths.split(','))
262    elif args.module:
263        print(f'Resolving test module {args.module}.')
264        for path in _get_module_artifacts(args.module):
265            if path.endswith(args.module):
266                mobly_bins.append(path)
267            if path.endswith('requirements.txt'):
268                requirements_files.append(path)
269            if path.endswith('.apk'):
270                test_apks.append(path)
271    elif args.packages or args.wheel:
272        packages = args.packages.split(',') if args.packages else [args.wheel]
273        unzip_root = tempfile.mkdtemp(prefix='mobly_unzip_')
274        _tempdirs.append(unzip_root)
275        for package in packages:
276            mobly_bins.append(os.path.abspath(package))
277            unzip_dir = os.path.join(unzip_root, os.path.basename(package))
278            print(f'Unzipping test package {package} to {unzip_dir}.')
279            os.makedirs(unzip_dir)
280            with zipfile.ZipFile(package) as zf:
281                zf.extractall(unzip_dir)
282            for root, _, files in os.walk(unzip_dir):
283                for file_name in files:
284                    path = os.path.join(root, file_name)
285                    if path.endswith('requirements.txt'):
286                        requirements_files.append(path)
287                    if path.endswith('.apk'):
288                        test_apks.append(path)
289    else:
290        print('No tests specified. Aborting.')
291        exit(1)
292    return mobly_bins, requirements_files, test_apks
293
294
295def _setup_virtualenv(
296        requirements_files: List[str],
297        wheel_file: Optional[str]
298) -> str:
299    """Creates a virtualenv and install dependencies into it.
300
301    Args:
302      requirements_files: List of paths of requirements.txt files.
303      wheel_file: A Mobly test package as an installable Python wheel.
304
305    Returns:
306      Path to the virtualenv's Python interpreter.
307    """
308    venv_dir = tempfile.mkdtemp(prefix='venv_')
309    _padded_print(f'Setting up virtualenv at {venv_dir}.')
310    subprocess.check_call([sys.executable, '-m', 'venv', venv_dir])
311    _tempdirs.append(venv_dir)
312    if platform.system() == 'Windows':
313        venv_executable = os.path.join(venv_dir, 'Scripts', 'python.exe')
314    else:
315        venv_executable = os.path.join(venv_dir, 'bin', 'python3')
316
317    # Install requirements
318    for requirements_file in requirements_files:
319        print(f'Installing dependencies from {requirements_file}.\n')
320        subprocess.check_call(
321            [venv_executable, '-m', 'pip', 'install', '-r', requirements_file]
322        )
323
324    # Install wheel
325    if wheel_file is not None:
326        print(f'Installing test wheel package {wheel_file}.\n')
327        subprocess.check_call(
328            [venv_executable, '-m', 'pip', 'install', wheel_file]
329        )
330    return venv_executable
331
332
333def _parse_adb_devices(lines: List[str]) -> List[str]:
334    """Parses result from 'adb devices' into a list of serials.
335
336    Derived from mobly.controllers.android_device.
337    """
338    results = []
339    for line in lines:
340        tokens = line.strip().split('\t')
341        if len(tokens) == 2 and tokens[1] == 'device':
342            results.append(tokens[0])
343    return results
344
345
346def _install_apks(
347        apks: List[str],
348        serials: Optional[List[str]] = None,
349) -> None:
350    """Installs given APKS to specified devices.
351
352    If no serials specified, installs APKs on all attached devices.
353
354    Args:
355      apks: List of paths to APKs.
356      serials: List of device serials.
357    """
358    _padded_print('Installing test APKs.')
359    if not serials:
360        adb_devices_out = (
361            subprocess.check_output(
362                ['adb', 'devices']
363            ).decode('utf-8').strip().splitlines()
364        )
365        serials = _parse_adb_devices(adb_devices_out)
366    for apk in apks:
367        for serial in serials:
368            print(f'Installing {apk} on device {serial}.')
369            subprocess.check_call(
370                ['adb', '-s', serial, 'install', '-r', '-g', apk]
371            )
372
373
374def _generate_mobly_config(serials: Optional[List[str]] = None) -> str:
375    """Generates a Mobly config for the provided device serials.
376
377    If no serials specified, generate a wildcard config (test loads all attached
378    devices).
379
380    Args:
381      serials: List of device serials.
382
383    Returns:
384      Path to the generated config.
385    """
386    config = {
387        'TestBeds': [{
388            'Name': _DEFAULT_TESTBED,
389            'Controllers': {
390                'AndroidDevice': serials if serials else '*',
391            },
392        }]
393    }
394    _, config_path = tempfile.mkstemp(prefix='mobly_config_', suffix='.yaml')
395    _padded_print(f'Generating Mobly config at {config_path}.')
396    with open(config_path, 'w') as f:
397        json.dump(config, f)
398    _tempfiles.append(config_path)
399    return config_path
400
401
402def _run_mobly_tests(
403        python_executable: Optional[str],
404        mobly_bins: List[str],
405        tests: Optional[List[str]],
406        config: str,
407        test_bed: str,
408        log_path: Optional[str]
409) -> None:
410    """Runs the Mobly tests with the specified binary and config."""
411    env = os.environ.copy()
412    base_log_path = _DEFAULT_MOBLY_LOGPATH
413    for mobly_bin in mobly_bins:
414        bin_name = os.path.basename(mobly_bin)
415        if log_path:
416            base_log_path = Path(log_path, bin_name)
417            env['MOBLY_LOGPATH'] = str(base_log_path)
418        cmd = [python_executable] if python_executable else []
419        cmd += [mobly_bin, '-c', config, '-tb', test_bed]
420        if tests is not None:
421            cmd.append('--tests')
422            cmd += tests
423        _padded_print(f'Running Mobly test {bin_name}.')
424        print(f'Command: {cmd}\n')
425        subprocess.run(cmd, env=env)
426        # Save a copy of the config in the log directory.
427        latest_logs = base_log_path.joinpath(test_bed, 'latest')
428        if latest_logs.is_dir():
429            shutil.copy2(config, latest_logs)
430
431
432def _clean_up() -> None:
433    """Cleans up temporary directories and files."""
434    _padded_print('Cleaning up temporary directories/files.')
435    for td in _tempdirs:
436        shutil.rmtree(td, ignore_errors=True)
437    _tempdirs.clear()
438    for tf in _tempfiles:
439        os.remove(tf)
440    _tempfiles.clear()
441
442
443def main() -> None:
444    args = _parse_args()
445
446    # args.module is not supported in Windows
447    if args.module and platform.system() == 'Windows':
448        print('The --module option is not supported in Windows. Aborting.')
449        exit(1)
450
451    # Build the test module if requested by user
452    if args.build:
453        _build_module(args.module)
454
455    serials = args.serials.split(',') if args.serials else None
456
457    # Extract test resources
458    mobly_bins, requirements_files, test_apks = _extract_test_resources(args)
459
460    # Install test APKs, if necessary
461    if args.install_apks:
462        _install_apks(test_apks, serials)
463
464    # Set up the Python virtualenv, if necessary
465    python_executable = None
466    if args.novenv:
467        if args.test_paths is not None:
468            python_executable = sys.executable
469    else:
470        python_executable = _setup_virtualenv(requirements_files, args.wheel)
471
472    if args.wheel:
473        mobly_bins = [
474            os.path.join(os.path.dirname(python_executable), args.bin)
475        ]
476        python_executable = None
477
478    # Generate the Mobly config, if necessary
479    config = args.config or _generate_mobly_config(serials)
480
481    # Run the tests
482    _run_mobly_tests(python_executable, mobly_bins, args.tests, config,
483                     args.test_bed, args.log_path)
484
485    # Clean up temporary dirs/files
486    _clean_up()
487
488
489if __name__ == '__main__':
490    main()
491