From 7c07b3957e7756610b0195908cb01ec20a80125a Mon Sep 17 00:00:00 2001
From: Alex Martani <alexandre.martani@benchling.com>
Date: Thu, 9 Nov 2023 15:01:10 -0800
Subject: [PATCH] feat: Test sharding support (#9)

---
 README.md                    |  3 +++
 e2e/smoke/.bazelrc           |  2 ++
 e2e/smoke/BUILD.bazel        | 11 +++++++++++
 e2e/smoke/requirements.in    |  1 +
 e2e/smoke/requirements.txt   |  6 ++++++
 e2e/smoke/test_sharded.py    |  9 +++++++++
 python_pytest/pytest_shim.py |  7 +++++++
 7 files changed, 39 insertions(+)
 create mode 100644 e2e/smoke/test_sharded.py

diff --git a/README.md b/README.md
index f4b05a2..0e4c72b 100644
--- a/README.md
+++ b/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
diff --git a/e2e/smoke/.bazelrc b/e2e/smoke/.bazelrc
index e69de29..af611e6 100644
--- a/e2e/smoke/.bazelrc
+++ b/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
diff --git a/e2e/smoke/BUILD.bazel b/e2e/smoke/BUILD.bazel
index 6620a64..693ef3b 100644
--- a/e2e/smoke/BUILD.bazel
+++ b/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 = [
diff --git a/e2e/smoke/requirements.in b/e2e/smoke/requirements.in
index 5662849..b74b182 100644
--- a/e2e/smoke/requirements.in
+++ b/e2e/smoke/requirements.in
@@ -1 +1,2 @@
 pytest==7.3.2
+pytest-shard==0.1.2
diff --git a/e2e/smoke/requirements.txt b/e2e/smoke/requirements.txt
index 60a6a11..1594aa8 100644
--- a/e2e/smoke/requirements.txt
+++ b/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 \
diff --git a/e2e/smoke/test_sharded.py b/e2e/smoke/test_sharded.py
new file mode 100644
index 0000000..2340fe1
--- /dev/null
+++ b/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"
diff --git a/python_pytest/pytest_shim.py b/python_pytest/pytest_shim.py
index 4041232..2ef54df 100644
--- a/python_pytest/pytest_shim.py
+++ b/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"):