|
| 1 | +"""A shim for executing pytest that supports test filtering, sharding, and more.""" |
| 2 | + |
| 3 | +import sys |
| 4 | +import os |
| 5 | + |
| 6 | +import pytest |
| 7 | + |
| 8 | + |
| 9 | +if __name__ == "__main__": |
| 10 | + pytest_args = ["--ignore=external"] |
| 11 | + |
| 12 | + args = sys.argv[1:] |
| 13 | + # pytest < 8.0 runs tests twice if __init__.py is passed explicitly as an argument. |
| 14 | + # Remove any __init__.py file to avoid that. |
| 15 | + # pytest.version_tuple is available since pytest 7.0 |
| 16 | + # https://github.com/pytest-dev/pytest/issues/9313 |
| 17 | + if not hasattr(pytest, "version_tuple") or pytest.version_tuple < (8, 0): |
| 18 | + args = [arg for arg in args if arg.startswith("-") or os.path.basename(arg) != "__init__.py"] |
| 19 | + |
| 20 | + if os.environ.get("XML_OUTPUT_FILE"): |
| 21 | + pytest_args.append("--junitxml={xml_output_file}".format(xml_output_file=os.environ.get("XML_OUTPUT_FILE"))) |
| 22 | + |
| 23 | + # Handle test sharding - requires pytest-shard plugin. |
| 24 | + if os.environ.get("TEST_SHARD_INDEX") and os.environ.get("TEST_TOTAL_SHARDS"): |
| 25 | + pytest_args.append("--shard-id={shard_id}".format(shard_id=os.environ.get("TEST_SHARD_INDEX"))) |
| 26 | + pytest_args.append("--num-shards={num_shards}".format(num_shards=os.environ.get("TEST_TOTAL_SHARDS"))) |
| 27 | + if os.environ.get("TEST_SHARD_STATUS_FILE"): |
| 28 | + open(os.environ["TEST_SHARD_STATUS_FILE"], "a").close() |
| 29 | + |
| 30 | + # Handle plugins that generate reports - if they are provided with relative paths (via args), |
| 31 | + # re-write it under bazel's test undeclared outputs dir. |
| 32 | + if os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR"): |
| 33 | + undeclared_output_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR") |
| 34 | + |
| 35 | + # Flags that take file paths as value. |
| 36 | + path_flags = [ |
| 37 | + "--report-log", # pytest-reportlog |
| 38 | + "--json-report-file", # pytest-json-report |
| 39 | + "--html", # pytest-html |
| 40 | + ] |
| 41 | + for i, arg in enumerate(args): |
| 42 | + for flag in path_flags: |
| 43 | + if arg.startswith(f"{flag}="): |
| 44 | + arg_split = arg.split("=", 1) |
| 45 | + if len(arg_split) == 2 and not os.path.isabs(arg_split[1]): |
| 46 | + args[i] = f"{flag}={undeclared_output_dir}/{arg_split[1]}" |
| 47 | + |
| 48 | + if os.environ.get("TESTBRIDGE_TEST_ONLY"): |
| 49 | + test_filter = os.environ["TESTBRIDGE_TEST_ONLY"] |
| 50 | + |
| 51 | + # If the test filter does not start with a class-like name, then use test filtering instead |
| 52 | + if not test_filter[0].isupper(): |
| 53 | + # --test_filter=test_module.test_fn or --test_filter=test_module/test_file.py |
| 54 | + pytest_args.extend(args) |
| 55 | + pytest_args.append("-k={filter}".format(filter=test_filter)) |
| 56 | + else: |
| 57 | + # --test_filter=TestClass.test_fn |
| 58 | + for arg in args: |
| 59 | + if not arg.startswith("--"): |
| 60 | + # arg is a src file. Add test class/method selection to it. |
| 61 | + # test.py::TestClass::test_fn |
| 62 | + arg = "{arg}::{module_fn}".format(arg=arg, module_fn=test_filter.replace(".", "::")) |
| 63 | + pytest_args.append(arg) |
| 64 | + else: |
| 65 | + pytest_args.extend(args) |
| 66 | + |
| 67 | + print(pytest_args, file=sys.stderr) |
| 68 | + raise SystemExit(pytest.main(pytest_args)) |
0 commit comments