55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
"""Tests for the streaming CLI runner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from swe_data_processing.cli import _run_streaming_stage
|
|
|
|
|
|
def test_streaming_stage_processes_records_concurrently(tmp_path: Path) -> None:
|
|
"""Multiple workers run processors in parallel while one thread writes JSONL."""
|
|
|
|
input_path = tmp_path / "input.jsonl"
|
|
output_path = tmp_path / "output.jsonl"
|
|
input_path.write_text(
|
|
"".join(
|
|
json.dumps({"trajectory_id": f"sample-{index}"}) + "\n" for index in range(20)
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
lock = threading.Lock()
|
|
active = 0
|
|
peak_active = 0
|
|
|
|
def processor(record: dict) -> dict:
|
|
nonlocal active, peak_active
|
|
with lock:
|
|
active += 1
|
|
peak_active = max(peak_active, active)
|
|
time.sleep(0.01)
|
|
with lock:
|
|
active -= 1
|
|
return {"sample_id": record["trajectory_id"]}
|
|
|
|
status = _run_streaming_stage(
|
|
input_path=input_path,
|
|
output_path=output_path,
|
|
errors_path=None,
|
|
limit=None,
|
|
resume=False,
|
|
workers=4,
|
|
stage_name="test",
|
|
processor=processor,
|
|
)
|
|
|
|
written = [json.loads(line) for line in output_path.read_text().splitlines()]
|
|
assert status == 0
|
|
assert peak_active == 4
|
|
assert {value["sample_id"] for value in written} == {
|
|
f"sample-{index}" for index in range(20)
|
|
}
|