Skip to content

Commit 78bec67

Browse files
committed
Fix KeyError when replacement worker has mismatched collection
When a worker crashes and is replaced, if the replacement worker's collection doesn't match the original, add_node_collection() returns early without adding the node to registered_collections. Later, when schedule() calls _reschedule() on all nodes, _assign_work_unit() crashes with KeyError accessing registered_collections[node]. Add a guard in _reschedule() to skip nodes that aren't in registered_collections. Fixes #1189 Fixes #714
1 parent 04c3ca8 commit 78bec67

2 files changed

Lines changed: 102 additions & 0 deletions

File tree

src/xdist/scheduler/loadscope.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,11 @@ def _reschedule(self, node: WorkerController) -> None:
320320
if node.shutting_down:
321321
return
322322

323+
# Skip nodes that aren't properly registered (e.g., collection mismatch)
324+
# See https://github.com/pytest-dev/pytest-xdist/issues/1189
325+
if node not in self.registered_collections:
326+
return
327+
323328
# Check that more work is available
324329
if not self.workqueue:
325330
node.shutdown()

testing/test_loadscope.py

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
"""Tests for LoadScopeScheduling."""
2+
3+
from __future__ import annotations
4+
5+
from collections.abc import Sequence
6+
from typing import TYPE_CHECKING
7+
8+
import execnet
9+
import pytest
10+
11+
from xdist.scheduler import LoadScopeScheduling
12+
from xdist.workermanage import WorkerController
13+
14+
15+
if TYPE_CHECKING:
16+
BaseOfMockGateway = execnet.Gateway
17+
BaseOfMockNode = WorkerController
18+
else:
19+
BaseOfMockGateway = object
20+
BaseOfMockNode = object
21+
22+
23+
class MockGateway(BaseOfMockGateway):
24+
_count = 0
25+
26+
def __init__(self) -> None:
27+
self.id = str(MockGateway._count)
28+
MockGateway._count += 1
29+
30+
31+
class MockNode(BaseOfMockNode):
32+
def __init__(self) -> None:
33+
self.sent: list[int] = []
34+
self.gateway = MockGateway()
35+
self._shutdown = False
36+
37+
def send_runtest_some(self, indices: Sequence[int]) -> None:
38+
self.sent.extend(indices)
39+
40+
def shutdown(self) -> None:
41+
self._shutdown = True
42+
43+
@property
44+
def shutting_down(self) -> bool:
45+
return self._shutdown
46+
47+
48+
@pytest.fixture(autouse=True)
49+
def reset_mock_gateway_counter() -> None:
50+
MockGateway._count = 0
51+
52+
53+
class TestLoadScopeScheduling:
54+
def test_replacement_worker_with_mismatched_collection_is_skipped(
55+
self, pytester: pytest.Pytester
56+
) -> None:
57+
"""Regression test for https://github.com/pytest-dev/pytest-xdist/issues/1189"""
58+
config = pytester.parseconfig("--tx=2*popen")
59+
sched = LoadScopeScheduling(config)
60+
61+
node1, node2 = MockNode(), MockNode()
62+
sched.add_node(node1)
63+
sched.add_node(node2)
64+
65+
collection = [
66+
"test_mod.py::test_a",
67+
"test_mod.py::test_b",
68+
"test_other.py::test_c",
69+
"test_other.py::test_d",
70+
]
71+
sched.add_node_collection(node1, collection)
72+
sched.add_node_collection(node2, collection)
73+
sched.schedule()
74+
75+
# Simulate node1 crashing
76+
sched.remove_node(node1)
77+
78+
# Replacement worker collects different tests (e.g., due to test file changes)
79+
replacement_node = MockNode()
80+
sched.add_node(replacement_node)
81+
different_collection = [
82+
"test_mod.py::test_a",
83+
"test_mod.py::test_b",
84+
"test_mod.py::test_NEW", # Different
85+
"test_other.py::test_d",
86+
]
87+
sched.add_node_collection(replacement_node, different_collection)
88+
89+
# Replacement node should not be in registered_collections due to mismatch
90+
assert replacement_node not in sched.registered_collections
91+
assert replacement_node in sched.assigned_work
92+
93+
# schedule() should skip unregistered nodes rather than crashing
94+
sched.schedule()
95+
96+
assert replacement_node.sent == []
97+
assert node2 in sched.registered_collections

0 commit comments

Comments
 (0)