xref: /aosp_15_r20/external/crosvm/tools/nextest_package (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1#!/usr/bin/env python3
2# Copyright 2023 The ChromiumOS Authors
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import argparse
7import json
8from multiprocessing.pool import ThreadPool
9import shlex
10import shutil
11from fnmatch import fnmatch
12from pathlib import Path
13from typing import Any, List, Tuple
14from impl.common import (
15    CROSVM_ROOT,
16    all_tracked_files,
17    chdir,
18    cmd,
19    cwd_context,
20    parallel,
21    print_timing_info,
22    quoted,
23    record_time,
24)
25
26# List of globs matching files in the source tree required by tests at runtime.
27# This is hard-coded specifically for crosvm tests.
28TEST_DATA_FILES = [
29    # Requried by nextest to obtain metadata
30    "*.toml",
31    # Configured by .cargo/config.toml to execute tests with the right emulator
32    ".cargo/runner.py",
33    # Requried by plugin tests
34    "crosvm_plugin/crosvm.h",
35    "tests/plugin.policy",
36]
37
38TEST_DATA_EXCLUDE = [
39    # config.toml is configured for x86 hosts. We cannot use that for remote tests.
40    ".cargo/config.toml",
41]
42
43cargo = cmd("cargo")
44tar = cmd("tar")
45rust_strip = cmd("rust-strip")
46
47
48def collect_rust_libs():
49    "Collect rust shared libraries required by the tests at runtime."
50    lib_dir = Path(cmd("rustc --print=sysroot").stdout()) / "lib"
51    for lib_file in lib_dir.glob("libstd-*"):
52        yield (lib_file, Path("debug/deps") / lib_file.name)
53    for lib_file in lib_dir.glob("libtest-*"):
54        yield (lib_file, Path("debug/deps") / lib_file.name)
55
56
57def collect_test_binaries(metadata: Any, strip: bool):
58    "Collect all test binaries that are needed to run the tests."
59    target_dir = Path(metadata["rust-build-meta"]["target-directory"])
60    test_binaries = [
61        Path(suite["binary-path"]).relative_to(target_dir)
62        for suite in metadata["rust-binaries"].values()
63    ]
64
65    non_test_binaries = [
66        Path(binary["path"])
67        for crate in metadata["rust-build-meta"]["non-test-binaries"].values()
68        for binary in crate
69    ]
70
71    def process_binary(binary_path: Path):
72        source_path = target_dir / binary_path
73        destination_path = binary_path
74        if strip:
75            stripped_path = source_path.with_suffix(".stripped")
76            if (
77                not stripped_path.exists()
78                or source_path.stat().st_ctime > stripped_path.stat().st_ctime
79            ):
80                rust_strip(f"--strip-all {source_path} -o {stripped_path}").fg()
81            return (stripped_path, destination_path)
82        else:
83            return (source_path, destination_path)
84
85    # Parallelize rust_strip calls.
86    pool = ThreadPool()
87    return pool.map(process_binary, test_binaries + non_test_binaries)
88
89
90def collect_test_data_files():
91    "List additional files from the source tree that are required by tests at runtime."
92    for file in all_tracked_files():
93        for glob in TEST_DATA_FILES:
94            if fnmatch(str(file), glob):
95                if str(file) not in TEST_DATA_EXCLUDE:
96                    yield (file, file)
97                break
98
99
100def collect_files(metadata: Any, output_directory: Path, strip_binaries: bool):
101    # List all files we need as (source path, path in output_directory) tuples
102    manifest: List[Tuple[Path, Path]] = [
103        *collect_test_binaries(metadata, strip=strip_binaries),
104        *collect_rust_libs(),
105        *collect_test_data_files(),
106    ]
107
108    # Create all target directories
109    for folder in set((output_directory / d).parent.resolve() for _, d in manifest):
110        folder.mkdir(exist_ok=True, parents=True)
111
112    # Use multiple processes of rsync copy the files fast (and only if needed)
113    parallel(
114        *(
115            cmd("rsync -a", source, output_directory / destination)
116            for source, destination in manifest
117        )
118    ).fg()
119
120
121def generate_run_script(metadata: Any, output_directory: Path):
122    # Generate metadata files for nextest
123    binares_metadata_file = "binaries-metadata.json"
124    (output_directory / binares_metadata_file).write_text(json.dumps(metadata))
125    cargo_metadata_file = "cargo-metadata.json"
126    cargo("metadata --format-version 1").write_to(output_directory / cargo_metadata_file)
127
128    # Put together command line to run nextest
129    run_cmd = [
130        "cargo-nextest",
131        "nextest",
132        "run",
133        f"--binaries-metadata={binares_metadata_file}",
134        f"--cargo-metadata={cargo_metadata_file}",
135        "--target-dir-remap=.",
136        "--workspace-remap=.",
137    ]
138    command_line = [
139        "#!/usr/bin/env bash",
140        'cd "$(dirname "${BASH_SOURCE[0]}")" || die',
141        f'{shlex.join(run_cmd)} "$@"',
142    ]
143
144    # Write command to a unix shell script
145    shell_script = output_directory / "run.sh"
146    shell_script.write_text("\n".join(command_line))
147    shell_script.chmod(0o755)
148
149    # TODO(denniskempin): Add an equivalent windows bash script
150
151
152def generate_archive(output_directory: Path, output_archive: Path):
153    with cwd_context(output_directory.parent):
154        tar("-ca", output_directory.name, "-f", output_archive).fg()
155
156
157def main():
158    """
159    Builds a package to execute tests remotely.
160
161    ## Basic usage
162
163    ```
164    $ tools/nextest_package -o foo.tar.zst ... nextest args
165    ```
166
167    The archive will contain all necessary test binaries, required shared libraries and test data
168    files required at runtime.
169    A cargo nextest binary is included along with a `run.sh` script to invoke it with the required
170    arguments. THe archive can be copied anywhere and executed:
171
172    ```
173    $ tar xaf foo.tar.zst && cd foo.tar.d && ./run.sh
174    ```
175
176    ## Nextest Arguments
177
178    All additional arguments will be passed to `nextest list`. Additional arguments to `nextest run`
179    can be passed to the `run.sh` invocation.
180
181    For example:
182
183    ```
184    $ tools/nextest_package -d foo --tests
185    $ cd foo && ./run.sh --test-threads=1
186    ```
187
188    Will only list and package integration tests (--tests) and run them with --test-threads=1.
189
190    ## Stripping Symbols
191
192    Debug symbols are stripped by default to reduce the package size. This can be disabled via
193    the `--no-strip` argument.
194
195    """
196    parser = argparse.ArgumentParser()
197    parser.add_argument("--no-strip", action="store_true")
198    parser.add_argument("--output-directory", "-d")
199    parser.add_argument("--output-archive", "-o")
200    parser.add_argument("--clean", action="store_true")
201    parser.add_argument("--timing-info", action="store_true")
202    (args, nextest_list_args) = parser.parse_known_args()
203    chdir(CROSVM_ROOT)
204
205    # Determine output archive / directory
206    output_directory = Path(args.output_directory).resolve() if args.output_directory else None
207    output_archive = Path(args.output_archive).resolve() if args.output_archive else None
208    if not output_directory and output_archive:
209        output_directory = output_archive.with_suffix(".d")
210    if not output_directory:
211        print("Must specify either --output-directory or --output-archive")
212        return
213
214    if args.clean and output_directory.exists():
215        shutil.rmtree(output_directory)
216
217    with record_time("Listing tests"):
218        cargo(
219            "nextest list",
220            *(quoted(a) for a in nextest_list_args),
221        ).fg()
222    with record_time("Listing tests metadata"):
223        metadata = cargo(
224            "nextest list --list-type binaries-only --message-format json",
225            *(quoted(a) for a in nextest_list_args),
226        ).json()
227
228    with record_time("Collecting files"):
229        collect_files(metadata, output_directory, strip_binaries=not args.no_strip)
230        generate_run_script(metadata, output_directory)
231
232    if output_archive:
233        with record_time("Generating archive"):
234            generate_archive(output_directory, output_archive)
235
236    if args.timing_info:
237        print_timing_info()
238
239
240if __name__ == "__main__":
241    main()
242