feat: Test sharding support (#9)

Author: amartani

README.md

@@ -34,6 +34,9 @@ py_pytest_test(
 )
 ```
 
+To use [test sharding](https://bazel.build/reference/test-encyclopedia#test-sharding), also add `requirement("pytest-shard")`
+to `deps`.
+
 ## Disclaimer
 
 I rely on this myself and will accept issue reports and contributions, however this only has a simple smoke test and isn't

e2e/smoke/.bazelrc

@@ -0,0 +1,2 @@
+# Enable after bazel 5.4.0 is removed from e2e tests to verify correct implementation of sharding.
+# test --incompatible_check_sharding_support

e2e/smoke/BUILD.bazel

@@ -15,6 +15,17 @@ py_pytest_test(
     ],
 )
 
+py_pytest_test(
+    name = "test_sharded",
+    size = "small",
+    srcs = ["test_sharded.py"],
+    shard_count = 2,
+    deps = [
+        requirement("pytest"),
+        requirement("pytest-shard"),
+    ],
+)
+
 build_test(
     name = "smoke_test",
     targets = [

e2e/smoke/requirements.in

@@ -1 +1,2 @@
 pytest==7.3.2
+pytest-shard==0.1.2

e2e/smoke/requirements.txt

@@ -23,6 +23,12 @@ pluggy==1.2.0 \
 pytest==7.3.2 \
     --hash=sha256:cdcbd012c9312258922f8cd3f1b62a6580fdced17db6014896053d47cddf9295 \
     --hash=sha256:ee990a3cc55ba808b80795a79944756f315c67c12b56abd3ac993a7b8c17030b
+    # via
+    #   -r requirements.in
+    #   pytest-shard
+pytest-shard==0.1.2 \
+    --hash=sha256:407a1df385cebe1feb9b4d2e7eeee8b044f8a24f0919421233159a17c59be2b9 \
+    --hash=sha256:b86a967fbfd1c8e50295095ccda031b7e890862ee06531d5142844f4c1d1cd67
     # via -r requirements.in
 tomli==2.0.1 \
     --hash=sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc \

e2e/smoke/test_sharded.py

@@ -0,0 +1,9 @@
+import os
+
+
+def test_shard_0():
+    assert os.environ["TEST_SHARD_INDEX"] == "0"
+
+
+def test_shard_1():
+    assert os.environ["TEST_SHARD_INDEX"] == "1"

python_pytest/pytest_shim.py

@@ -16,6 +16,13 @@
     if os.environ.get("XML_OUTPUT_FILE"):
         pytest_args.append("--junitxml={xml_output_file}".format(xml_output_file=os.environ.get("XML_OUTPUT_FILE")))
 
+    # Handle test sharding - requires pytest-shard plugin.
+    if os.environ.get("TEST_SHARD_INDEX") and os.environ.get("TEST_TOTAL_SHARDS"):
+        pytest_args.append("--shard-id={shard_id}".format(shard_id=os.environ.get("TEST_SHARD_INDEX")))
+        pytest_args.append("--num-shards={num_shards}".format(num_shards=os.environ.get("TEST_TOTAL_SHARDS")))
+        if os.environ.get("TEST_SHARD_STATUS_FILE"):
+            open(os.environ["TEST_SHARD_STATUS_FILE"], "a").close()
+
     # Handle plugins that generate reports - if they are provided with relative paths (via args),
     # re-write it under bazel's test undeclared outputs dir.
     if os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR"):