xref: /aosp_15_r20/external/bazelbuild-rules_python/tests/integration/runner.py (revision 60517a1edbc8ecf509223e9af94a7adec7d736b8)
1# Copyright 2024 The Bazel Authors. All rights reserved.
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
15import logging
16import os
17import os.path
18import pathlib
19import re
20import shlex
21import subprocess
22import unittest
23
24_logger = logging.getLogger(__name__)
25
26class ExecuteError(Exception):
27    def __init__(self, result):
28        self.result = result
29    def __str__(self):
30        return self.result.describe()
31
32class ExecuteResult:
33    def __init__(
34        self,
35        args: list[str],
36        env: dict[str, str],
37        cwd: pathlib.Path,
38        proc_result: subprocess.CompletedProcess,
39    ):
40        self.args = args
41        self.env = env
42        self.cwd = cwd
43        self.exit_code = proc_result.returncode
44        self.stdout = proc_result.stdout
45        self.stderr = proc_result.stderr
46
47    def describe(self) -> str:
48        env_lines = [
49            "  " + shlex.quote(f"{key}={value}")
50            for key, value in sorted(self.env.items())
51        ]
52        env = " \\\n".join(env_lines)
53        args = shlex.join(self.args)
54        maybe_stdout_nl = "" if self.stdout.endswith("\n") else "\n"
55        maybe_stderr_nl = "" if self.stderr.endswith("\n") else "\n"
56        return f"""\
57COMMAND:
58cd {self.cwd} && \\
59env \\
60{env} \\
61  {args}
62RESULT: exit_code: {self.exit_code}
63===== STDOUT START =====
64{self.stdout}{maybe_stdout_nl}===== STDOUT END   =====
65===== STDERR START =====
66{self.stderr}{maybe_stderr_nl}===== STDERR END   =====
67"""
68
69
70class TestCase(unittest.TestCase):
71    def setUp(self):
72        super().setUp()
73        self.repo_root = pathlib.Path(os.environ["BIT_WORKSPACE_DIR"])
74        self.bazel = pathlib.Path(os.environ["BIT_BAZEL_BINARY"])
75        outer_test_tmpdir = pathlib.Path(os.environ["TEST_TMPDIR"])
76        self.test_tmp_dir = outer_test_tmpdir / "bit_test_tmp"
77        # Put the global tmp not under the test tmp to better match how a real
78        # execution has entirely different directories for these.
79        self.tmp_dir = outer_test_tmpdir / "bit_tmp"
80        self.bazel_env = {
81            "PATH": os.environ["PATH"],
82            "TEST_TMPDIR": str(self.test_tmp_dir),
83            "TMP": str(self.tmp_dir),
84            # For some reason, this is necessary for Bazel 6.4 to work.
85            # If not present, it can't find some bash helpers in @bazel_tools
86            "RUNFILES_DIR": os.environ["TEST_SRCDIR"]
87        }
88
89    def run_bazel(self, *args: str, check: bool = True) -> ExecuteResult:
90        """Run a bazel invocation.
91
92        Args:
93            *args: The args to pass to bazel; the leading `bazel` command is
94                added automatically
95            check: True if the execution must succeed, False if failure
96                should raise an error.
97        Returns:
98            An `ExecuteResult` from running Bazel
99        """
100        args = [str(self.bazel), *args]
101        env = self.bazel_env
102        _logger.info("executing: %s", shlex.join(args))
103        cwd = self.repo_root
104        proc_result = subprocess.run(
105            args=args,
106            text=True,
107            capture_output=True,
108            cwd=cwd,
109            env=env,
110            check=False,
111        )
112        exec_result = ExecuteResult(args, env, cwd, proc_result)
113        if check and exec_result.exit_code:
114            raise ExecuteError(exec_result)
115        else:
116            return exec_result
117
118    def assert_result_matches(self, result: ExecuteResult, regex: str) -> None:
119        """Assert stdout/stderr of an invocation matches a regex.
120
121        Args:
122            result: ExecuteResult from `run_bazel` whose stdout/stderr will
123                be checked.
124            regex: Pattern to match, using `re.search` semantics.
125        """
126        if not re.search(regex, result.stdout + result.stderr):
127            self.fail(
128                "Bazel output did not match expected pattern\n"
129                + f"expected pattern: {regex}\n"
130                + f"invocation details:\n{result.describe()}"
131            )
132