48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
"""Tests for deterministic review sampling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from swe_data_processing.sampling import iter_review_records
|
|
|
|
|
|
def test_review_reads_only_provenance_shards(tmp_path: Path) -> None:
|
|
selected_shard = tmp_path / "selected.parquet"
|
|
unrelated_shard = tmp_path / "unrelated.parquet"
|
|
pq.write_table(
|
|
pa.Table.from_pylist(
|
|
[
|
|
{"trajectory_id": "wanted", "trajectory": "selected"},
|
|
{"trajectory_id": "other", "trajectory": "same shard"},
|
|
]
|
|
),
|
|
selected_shard,
|
|
)
|
|
pq.write_table(
|
|
pa.Table.from_pylist(
|
|
[{"trajectory_id": "unrelated", "trajectory": "must not be read"}]
|
|
),
|
|
unrelated_shard,
|
|
)
|
|
selected = {"wanted": ["hard:FIVE_CONSECUTIVE_TOOL_FAILURES"]}
|
|
decision = {
|
|
"sample_id": "wanted",
|
|
"source_parquet": str(selected_shard),
|
|
"hard_reject_reasons": ["FIVE_CONSECUTIVE_TOOL_FAILURES"],
|
|
}
|
|
|
|
records = list(
|
|
iter_review_records(
|
|
tmp_path,
|
|
selected,
|
|
{"wanted": decision},
|
|
)
|
|
)
|
|
|
|
assert [record["trajectory_id"] for record in records] == ["wanted"]
|
|
assert records[0]["_qc_review"]["decision"] == decision
|