test(tests): stabilize write_json concurrent serialization flake

This commit is contained in:
fangliquanflq
2026-08-08 04:24:35 +00:00
committed by Teknium
parent 7c2bc87f81
commit d3fca90fa6
+52 -11
View File
@@ -896,24 +896,65 @@ class _BrokenStdout:
def test_write_json_serializes_concurrent_writes(monkeypatch):
out = _ChunkyStdout()
monkeypatch.setattr(server, "_real_stdout", out)
"""Assert StdioTransport holds _stdout_lock across the full stream.write.
threads = [
threading.Thread(target=server.write_json, args=({"seq": i, "text": "x" * 24},))
for i in range(8)
]
The old char-by-char sleep mock made this test take long enough that
leftover background write_json calls from earlier cases in this file
could append an extra line (intermittent ``assert 9 == 8`` on CI/main).
Match the WS concurrent-send check: count in-flight writes, and only
assert on frames that carry this test's marker payload.
"""
marker = "x" * 24
active = 0
max_active = 0
gate = threading.Lock()
frames: list[str] = []
class RecordingStdout:
def write(self, text: str) -> int:
nonlocal active, max_active
with gate:
active += 1
max_active = max(max_active, active)
try:
# Release the GIL while "in write" so a missing outer lock
# would let another thread bump max_active above 1.
time.sleep(0.01)
frames.append(text)
finally:
with gate:
active -= 1
return len(text)
def flush(self) -> None:
return None
monkeypatch.setattr(server, "_real_stdout", RecordingStdout())
barrier = threading.Barrier(8)
def _worker(seq: int) -> None:
barrier.wait(timeout=5)
server.write_json({"seq": seq, "text": marker})
threads = [threading.Thread(target=_worker, args=(i,)) for i in range(8)]
for t in threads:
t.start()
for t in threads:
t.join()
t.join(timeout=10)
assert not t.is_alive()
lines = "".join(out.parts).splitlines()
assert max_active == 1
assert len(lines) == 8
assert {json.loads(line)["seq"] for line in lines} == set(range(8))
ours = []
for frame in frames:
assert frame.endswith("\n"), frame
obj = json.loads(frame)
if obj.get("text") == marker and "seq" in obj:
ours.append(obj)
assert {obj["seq"] for obj in ours} == set(range(8))
assert len(ours) == 8
def test_write_json_returns_false_on_broken_pipe(monkeypatch):