File size: 4,083 Bytes
9a46c96 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | 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)
|