• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# Copyright 2021 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file
5import argparse
6import platform
7import subprocess
8from pathlib import Path
9from typing import Any, Literal, Optional, cast
10import typing
11import sys
12import testvm
13import os
14
15USAGE = """Choose to run tests locally, in a vm or on a remote machine.
16
17To set the default test target to run on one of the build-in VMs:
18
19    ./tools/set_test_target vm:aarch64 && source .envrc
20
21Then as usual run cargo or run_tests:
22
23    ./tools/run_tests
24    cargo test
25
26The command will prepare the VM for testing (e.g. upload required shared
27libraries for running rust tests) and set up run_tests as well as cargo
28to build for the test target and execute tests on it.
29
30Arbitrary SSH remotes can be used for running tests as well. e.g.
31
32    ./tools/set_test_target ssh:remotehost
33
34The `remotehost` needs to be properly configured for passwordless
35authentication.
36
37Tip: Use http://direnv.net to automatically load the envrc file instead of
38having to source it after each call.
39"""
40
41SCRIPT_PATH = Path(__file__).resolve()
42SCRIPT_DIR = SCRIPT_PATH.parent.resolve()
43TESTVM_DIR = SCRIPT_DIR.parent.joinpath("testvm")
44TARGET_DIR = testvm.cargo_target_dir().joinpath("crosvm_tools")
45ENVRC_PATH = SCRIPT_DIR.parent.parent.joinpath(".envrc")
46
47Arch = Literal["x86_64", "aarch64", "armhf", "win64"]
48
49# Enviroment variables needed for building with cargo
50BUILD_ENV = {
51    "PKG_CONFIG_aarch64_unknown_linux_gnu": "aarch64-linux-gnu-pkg-config",
52    "PKG_CONFIG_armv7_unknown_linux_gnueabihf": "arm-linux-gnueabihf-pkg-config",
53    "CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER": "aarch64-linux-gnu-gcc",
54}
55
56
57class Ssh:
58    """Wrapper around subprocess to execute commands remotely via SSH."""
59
60    hostname: str
61    opts: list[str]
62
63    def __init__(self, hostname: str, opts: list[str] = []):
64        self.hostname = hostname
65        self.opts = opts
66
67    def run(self, cmd: str, **kwargs: Any):
68        """Equivalent of subprocess.run"""
69        return subprocess.run(
70            [
71                "ssh",
72                self.hostname,
73                *self.opts,
74                # Do not create a tty. This will mess with terminal output
75                # when running multiple subprocesses.
76                "-T",
77                # Tell sh to kill children on hangup.
78                f"shopt -s huponexit; {cmd}",
79            ],
80            **kwargs,
81        )
82
83    def check_output(self, cmd: str):
84        """Equivalent of subprocess.check_output"""
85        return subprocess.run(
86            ["ssh", self.hostname, *self.opts, "-T", cmd],
87            stdout=subprocess.PIPE,
88            stderr=subprocess.STDOUT,
89            text=True,
90            check=True,
91        ).stdout
92
93    def upload_files(self, files: list[Path], remote_dir: str = "", quiet: bool = False):
94        """Wrapper around SCP."""
95        flags: list[str] = []
96        if quiet:
97            flags.append("-q")
98        scp_cmd = [
99            "scp",
100            *flags,
101            *self.opts,
102            *files,
103            f"{self.hostname}:{remote_dir}",
104        ]
105        subprocess.run(scp_cmd, check=True)
106
107
108class TestTarget(object):
109    """A test target can be the local host, a VM or a remote devica via SSH."""
110
111    target_str: str
112    is_host: bool = False
113    vm: Optional[testvm.Arch] = None
114    ssh: Optional[Ssh] = None
115    __arch: Optional[Arch] = None
116
117    @classmethod
118    def default(cls):
119        return cls(os.environ.get("CROSVM_TEST_TARGET", "host"))
120
121    def __init__(self, target_str: str):
122        """target_str can be "vm:arch", "ssh:hostname" or just "host" """
123        self.target_str = target_str
124        parts = target_str.split(":")
125        if len(parts) == 2 and parts[0] == "vm":
126            arch: testvm.Arch = parts[1]  # type: ignore
127            self.vm = arch
128            self.ssh = Ssh("localhost", testvm.ssh_cmd_args(arch))
129        elif len(parts) == 2 and parts[0] == "ssh":
130            self.ssh = Ssh(parts[1])
131        elif len(parts) == 1 and parts[0] == "host":
132            self.is_host = True
133        else:
134            raise Exception(f"Invalid target {target_str}")
135
136    @property
137    def arch(self) -> Arch:
138        if not self.__arch:
139            if self.vm:
140                self.__arch = self.vm
141            elif self.ssh:
142                self.__arch = cast(Arch, self.ssh.check_output("arch").strip())
143            else:
144                self.__arch = cast(Arch, platform.machine())
145        return self.__arch
146
147    def __str__(self):
148        return self.target_str
149
150
151def find_rust_lib_dir():
152    cargo_path = Path(subprocess.check_output(["rustup", "which", "cargo"], text=True))
153    if os.name == "posix":
154        return cargo_path.parent.parent.joinpath("lib")
155    elif os.name == "nt":
156        return cargo_path.parent
157    else:
158        raise Exception(f"Unsupported build target: {os.name}")
159
160
161def find_rust_libs():
162    lib_dir = find_rust_lib_dir()
163    yield from lib_dir.glob("libstd-*")
164    yield from lib_dir.glob("libtest-*")
165
166
167def prepare_remote(ssh: Ssh, extra_files: list[Path] = []):
168    print("Preparing remote")
169    ssh.upload_files(list(find_rust_libs()) + extra_files)
170    pass
171
172
173def prepare_target(target: TestTarget, extra_files: list[Path] = []):
174    if target.vm:
175        testvm.build_if_needed(target.vm)
176        testvm.wait(target.vm)
177    if target.ssh:
178        prepare_remote(target.ssh, extra_files)
179
180
181def get_cargo_build_target(arch: Arch):
182    if os.name == "posix":
183        if arch == "armhf":
184            return "armv7-unknown-linux-gnueabihf"
185        elif arch == "win64":
186            return "x86_64-pc-windows-gnu"
187        else:
188            return f"{arch}-unknown-linux-gnu"
189    elif os.name == "nt":
190        if arch == "win64":
191            return f"x86_64-pc-windows-msvc"
192        else:
193            return f"{arch}-pc-windows-msvc"
194    else:
195        raise Exception(f"Unsupported build target: {os.name}")
196
197
198def get_cargo_env(target: TestTarget, build_arch: Arch):
199    """Environment variables to make cargo use the test target."""
200    env: dict[str, str] = BUILD_ENV.copy()
201    cargo_target = get_cargo_build_target(build_arch)
202    upper_target = cargo_target.upper().replace("-", "_")
203    if build_arch != platform.machine():
204        env["CARGO_BUILD_TARGET"] = cargo_target
205    if not target.is_host:
206        env[f"CARGO_TARGET_{upper_target}_RUNNER"] = f"{SCRIPT_PATH} exec-file"
207    env["CROSVM_TEST_TARGET"] = str(target)
208    return env
209
210
211def write_envrc(values: dict[str, str]):
212    with open(ENVRC_PATH, "w") as file:
213        for key, value in values.items():
214            file.write(f'export {key}="{value}"\n')
215
216
217def set_target(target: TestTarget, build_arch: Optional[Arch]):
218    prepare_target(target)
219    if not build_arch:
220        build_arch = target.arch
221    write_envrc(get_cargo_env(target, build_arch))
222    print(f"Test target: {target}")
223    print(f"Target Architecture: {build_arch}")
224
225
226def exec_file_on_target(
227    target: TestTarget,
228    filepath: Path,
229    timeout: int,
230    args: list[str] = [],
231    extra_files: list[Path] = [],
232    **kwargs: Any,
233):
234    """Executes a file on the test target.
235
236    The file is uploaded to the target's home directory (if it's an ssh or vm
237    target) plus any additional extra files provided, then executed and
238    deleted afterwards.
239
240    If the test target is 'host', files will just be executed locally.
241
242    Timeouts will trigger a subprocess.TimeoutExpired exception, which contanins
243    any output produced by the subprocess until the timeout.
244    """
245    env = os.environ.copy()
246    if not target.ssh:
247        # Allow test binaries to find rust's test libs.
248        if os.name == "posix":
249            env["LD_LIBRARY_PATH"] = str(find_rust_lib_dir())
250        elif os.name == "nt":
251            if not env["PATH"]:
252                env["PATH"] = str(find_rust_lib_dir())
253            else:
254                env["PATH"] += ";" + str(find_rust_lib_dir())
255        else:
256            raise Exception(f"Unsupported build target: {os.name}")
257
258        cmd_line = [str(filepath), *args]
259        return subprocess.run(
260            cmd_line,
261            env=env,
262            timeout=timeout,
263            text=True,
264            **kwargs,
265        )
266    else:
267        filename = Path(filepath).name
268        target.ssh.upload_files([filepath] + extra_files, quiet=True)
269        try:
270            result = target.ssh.run(
271                f"chmod +x {filename} && sudo LD_LIBRARY_PATH=. ./{filename} {' '.join(args)}",
272                timeout=timeout,
273                text=True,
274                **kwargs,
275            )
276        finally:
277            # Remove uploaded files regardless of test result
278            all_filenames = [filename] + [f.name for f in extra_files]
279            target.ssh.check_output(f"sudo rm {' '.join(all_filenames)}")
280        return result
281
282
283def exec_file(
284    target: TestTarget,
285    filepath: Path,
286    args: list[str] = [],
287    timeout: int = 60,
288    extra_files: list[Path] = [],
289):
290    if not filepath.exists():
291        raise Exception(f"File does not exist: {filepath}")
292
293    print(f"Executing `{Path(filepath).name} {' '.join(args)}` on {target}")
294    try:
295        sys.exit(exec_file_on_target(target, filepath, timeout, args, extra_files).returncode)
296    except subprocess.TimeoutExpired as e:
297        print(f"Process timed out after {e.timeout}s")
298
299
300def main():
301    COMMANDS = [
302        "set",
303        "exec-file",
304    ]
305
306    parser = argparse.ArgumentParser(usage=USAGE)
307    parser.add_argument("command", choices=COMMANDS)
308    parser.add_argument("--target", type=str, help="Override default test target.")
309    parser.add_argument(
310        "--arch",
311        choices=typing.get_args(Arch),
312        help="Override target build architecture.",
313    )
314    parser.add_argument(
315        "--extra-files",
316        type=str,
317        nargs="*",
318        default=[],
319        help="Additional files required by the binary to execute.",
320    )
321    parser.add_argument(
322        "--timeout",
323        type=int,
324        default=60,
325        help="Kill the process after the specified timeout.",
326    )
327    parser.add_argument("remainder", nargs=argparse.REMAINDER)
328    args = parser.parse_args()
329
330    if args.command == "set":
331        if len(args.remainder) != 1:
332            parser.error("Need to specify a target.")
333        set_target(TestTarget(args.remainder[0]), args.arch)
334        return
335
336    if args.target:
337        target = TestTarget(args.target)
338    else:
339        target = TestTarget.default()
340
341    if args.command == "exec-file":
342        if len(args.remainder) < 1:
343            parser.error("Need to specify a file to execute.")
344        exec_file(
345            target,
346            Path(args.remainder[0]),
347            args=args.remainder[1:],
348            timeout=args.timeout,
349            extra_files=[Path(f) for f in args.extra_files],
350        )
351
352
353if __name__ == "__main__":
354    try:
355        main()
356    except subprocess.CalledProcessError as e:
357        print("Command failed:", e.cmd)
358        print(e.stdout)
359        print(e.stderr)
360        sys.exit(-1)
361