Skip to content

Commit

Permalink
[test_utilities] Obey test sharding protocol during Python tests (Rob…
Browse files Browse the repository at this point in the history
…otLocomotion#20660)

The lack of affirmation will become a console warning in Bazel 7.
  • Loading branch information
jwnimmer-tri authored Dec 12, 2023
1 parent 4fbbce6 commit 81d6caa
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions common/test_utilities/drake_py_unittest_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from importlib.machinery import SourceFileLoader
import io
import os
from pathlib import Path
import re
import sys
import trace
Expand All @@ -23,6 +24,11 @@

def _unittest_main(*, module, argv, testRunner):
"""Just like unittest.main, but obeys TEST_TOTAL_SHARDS.
Refer to https://bazel.build/reference/test-encyclopedia#test-sharding for
an overview of the sharding protocol used by the Bazel test runner that
invokes us.
Only a subset of unittest.main's kwargs are supported.
"""
# In case sharding won't be used, delegate to the vanilla unittest.main.
Expand All @@ -39,6 +45,10 @@ def _unittest_main(*, module, argv, testRunner):
testRunner=testRunner)
return

# Affirm to the test environment that we are cognizant of sharding.
if "TEST_SHARD_STATUS_FILE" in os.environ:
Path(os.environ["TEST_SHARD_STATUS_FILE"]).touch()

# Run only some of the tests, per the BUILD.bazel's shard_count.
total_shards = int(os.environ["TEST_TOTAL_SHARDS"])
shard_index = int(os.environ["TEST_SHARD_INDEX"])
Expand Down

0 comments on commit 81d6caa

Please sign in to comment.