| import json |
| from decimal import Decimal |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from dlb.io import SampleCountError, SampleValidationError, validate_samples, write_samples_atomic |
|
|
|
|
| def make_records(count: int) -> list[dict[str, object]]: |
| return [ |
| { |
| "sample_id": index, |
| "text": f"sample {index}", |
| "token_ids": [index + 1], |
| "seed": 42, |
| "generation_seconds": 0.25, |
| } |
| for index in range(count) |
| ] |
|
|
|
|
| def test_validate_requires_exact_count(tmp_path: Path) -> None: |
| """Catch a completed-looking file with the wrong requested sample count.""" |
|
|
| path = tmp_path / "samples.jsonl" |
| write_samples_atomic(path, make_records(3)) |
|
|
| with pytest.raises(SampleCountError, match="expected 1024 records, found 3"): |
| validate_samples(path, expected=1024) |
|
|
|
|
| def test_writer_keeps_existing_final_when_new_records_are_invalid(tmp_path: Path) -> None: |
| """Catch publication that replaces a valid final file before semantic validation.""" |
|
|
| path = tmp_path / "samples.jsonl" |
| write_samples_atomic(path, make_records(1)) |
|
|
| with pytest.raises(SampleValidationError, match="record 0") as error: |
| write_samples_atomic( |
| path, |
| [ |
| { |
| "sample_id": 0, |
| "text": " ", |
| "token_ids": [1], |
| "seed": 42, |
| "generation_seconds": 0.1, |
| } |
| ], |
| ) |
|
|
| assert "text" in str(error.value) |
|
|
| assert [json.loads(line) for line in path.read_text().splitlines()] == make_records(1) |
| assert not list(tmp_path.glob(".samples.jsonl.*.partial")) |
|
|
|
|
| def test_validator_reports_record_index_for_noncontiguous_ids(tmp_path: Path) -> None: |
| """Catch duplicate or missing sample identifiers hidden in otherwise valid JSON.""" |
|
|
| path = tmp_path / "samples.jsonl" |
| path.write_text("\n".join(json.dumps(item) for item in make_records(2)) + "\n") |
| rows = path.read_text().splitlines() |
| second = json.loads(rows[1]) |
| second["sample_id"] = 3 |
| path.write_text(rows[0] + "\n" + json.dumps(second) + "\n") |
|
|
| with pytest.raises(SampleValidationError, match="record 1.*expected sample_id 1, found 3"): |
| validate_samples(path, expected=2) |
|
|
|
|
| def test_writer_rejects_symlinked_final_target(tmp_path: Path) -> None: |
| """Catch a symlink attack that would redirect an atomic replace outside results.""" |
|
|
| target = tmp_path / "target.jsonl" |
| target.write_text("outside\n") |
| path = tmp_path / "samples.jsonl" |
| path.symlink_to(target) |
|
|
| with pytest.raises(SampleValidationError, match="symlink"): |
| write_samples_atomic(path, make_records(1)) |
|
|
| assert target.read_text() == "outside\n" |
|
|
|
|
| def test_writer_counts_one_record_from_a_generator(tmp_path: Path) -> None: |
| """Catch iterable publication that validates a one-row generator as empty.""" |
|
|
| path = tmp_path / "samples.jsonl" |
| count = write_samples_atomic(path, (record for record in make_records(1))) |
|
|
| assert count == 1 |
| assert validate_samples(path, expected=1) == 1 |
|
|
|
|
| @pytest.mark.parametrize("timing", [True, 1, Decimal("0.1"), "0.1", float("nan"), float("inf")]) |
| def test_writer_rejects_non_float_or_nonfinite_generation_time(tmp_path: Path, timing: object) -> None: |
| """Catch coercive timing values before they enter a published sample artifact.""" |
|
|
| record = make_records(1)[0] |
| record["generation_seconds"] = timing |
|
|
| with pytest.raises(SampleValidationError, match="generation_seconds"): |
| write_samples_atomic(tmp_path / "samples.jsonl", [record]) |
|
|
|
|
| def test_validator_rejects_integer_generation_time_from_json(tmp_path: Path) -> None: |
| """Catch a JSON numeric integer accepted as a timing float during streamed validation.""" |
|
|
| record = make_records(1)[0] |
| record["generation_seconds"] = 1 |
| path = tmp_path / "samples.jsonl" |
| path.write_text(json.dumps(record) + "\n") |
|
|
| with pytest.raises(SampleValidationError, match="generation_seconds"): |
| validate_samples(path, expected=1) |
|
|
