79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
"""Tests for CLI parsing and streaming profile execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
from swe_data_processing.cli import (
|
|
_run_parallel_parquet_profile,
|
|
_run_streaming_stage,
|
|
build_parser,
|
|
)
|
|
|
|
|
|
def test_streaming_stage_writes_each_record(tmp_path: Path) -> None:
|
|
input_path = tmp_path / "input.jsonl"
|
|
output_path = tmp_path / "output.jsonl"
|
|
input_path.write_text(
|
|
"".join(json.dumps({"trajectory_id": f"s-{i}"}) + "\n" for i in range(10)),
|
|
encoding="utf-8",
|
|
)
|
|
status = _run_streaming_stage(
|
|
input_path=input_path,
|
|
output_path=output_path,
|
|
errors_path=None,
|
|
limit=None,
|
|
resume=False,
|
|
workers=3,
|
|
processor=lambda record: {"sample_id": record["trajectory_id"]},
|
|
)
|
|
assert status == 0
|
|
assert len(output_path.read_text().splitlines()) == 10
|
|
|
|
|
|
def test_cli_has_only_deterministic_pipeline_commands() -> None:
|
|
parser = build_parser()
|
|
assert parser.parse_args(
|
|
["profile", "--input", "in.jsonl", "--output", "out.jsonl"]
|
|
).command == "profile"
|
|
help_text = parser.format_help()
|
|
assert "classify" not in help_text
|
|
assert "repair" not in help_text
|
|
assert "materialize" in help_text
|
|
|
|
|
|
def test_parallel_parquet_profile_uses_atomic_shard_parts(tmp_path: Path) -> None:
|
|
dataset = tmp_path / "dataset"
|
|
dataset.mkdir()
|
|
for index in range(2):
|
|
pq.write_table(
|
|
pa.Table.from_pylist(
|
|
[{"trajectory_id": f"sample-{index}", "resolved": 1}]
|
|
),
|
|
dataset / f"part-{index}.parquet",
|
|
)
|
|
output = tmp_path / "metrics.jsonl"
|
|
errors = tmp_path / "errors.jsonl"
|
|
status = _run_parallel_parquet_profile(
|
|
Namespace(
|
|
input=dataset,
|
|
output=output,
|
|
errors=errors,
|
|
tokenizer_json=None,
|
|
workers=2,
|
|
resume=False,
|
|
)
|
|
)
|
|
|
|
assert status == 0
|
|
assert [json.loads(line)["sample_id"] for line in output.read_text().splitlines()] == [
|
|
"sample-0",
|
|
"sample-1",
|
|
]
|
|
assert errors.read_text() == ""
|