Last active
February 14, 2024 10:24
-
-
Save andgineer/9679190630c7845023546ef24bda2d94 to your computer and use it in GitHub Desktop.
pytest-xdist custom scheduler: Executes tests from the priority list first. Placing the longest tests in this list could significantly reduce test suite run time.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def pytest_xdist_make_scheduler(config, log): | |
"""Create a custom scheduler for pytest-xdist.""" | |
return PriorityListScheduling(config, log) | |
"""Custom scheduler implementation that run tests from tests/resources/priority_tests.txt first. | |
For example that enables us to run long tests first and on separate nodes. | |
To get tests list sorted by duration we can use pytest --durations=200 | |
""" | |
from datetime import datetime | |
from xdist.dsession import LoadScheduling | |
from xdist.workermanage import WorkerController | |
class PriorityListScheduling(LoadScheduling): | |
_prioritized_tests_indices = None | |
def __init__(self, config, log): | |
super().__init__(config, log) | |
self.prioritized_tests = self.load_prioritized_tests() | |
self.debug(f"Custom xdist scheduler have loaded {len(self.prioritized_tests)} long tests.") | |
def debug(self, *message): | |
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
full_message = f"[#]{timestamp}[#] {' '.join(message)}" | |
print(full_message) | |
@property | |
def prioritized_tests_indices(self): | |
"""Map prioritized test names to indices. | |
At __init__ tests are not collected so we do lazy initialization. | |
Calculate at first access and use cache afterward. | |
""" | |
if self._prioritized_tests_indices is None: | |
self._prioritized_tests_indices = [ | |
self.collection.index(name) # we could create reverse-index but not worth it - less than 100 tests | |
for name in self.prioritized_tests | |
if name in self.collection | |
] | |
return self._prioritized_tests_indices | |
def _send_tests(self, node: WorkerController, num: int) -> None: | |
"""Called by xdist. | |
(!) todo research more stable solution. could break in future versions of xdist. | |
""" | |
# First check if there are prioritized tests to send | |
test_to_send = next( | |
( | |
test | |
for test in self.prioritized_tests_indices | |
if test in self.pending | |
) | |
) | |
if test_to_send: | |
self.debug(f"Send prioritized test {self.collection[test_to_send]} to the node {node.gateway.id}") | |
self.pending.remove(test_to_send) | |
self.node2pending[node].append(test_to_send) | |
node.send_runtest_some([test_to_send]) | |
return # If we've sent prioritized test we are done | |
# default logic of '--dist load' scheduler | |
# todo switch to '--dist loadfile'? | |
self.debug(f"Send {num} un-prioritized tests to the node {node.gateway.id}") | |
tests_per_node = self.pending[:num] | |
if tests_per_node: | |
del self.pending[:num] | |
self.node2pending[node].extend(tests_per_node) | |
node.send_runtest_some(tests_per_node) | |
def load_prioritized_tests(self, filename="tests/resources/priority_tests.txt"): | |
try: | |
with open(filename, "r") as f: | |
return [line.strip() for line in f if line.strip() if not line.startswith("#")] | |
except FileNotFoundError: | |
self.debug(f"Custom xdist scheduler not found '{filename}'. No tests will be prioritized.") | |
return [] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment