| """Regression tests for audit findings U09-1..U09-4 (align/data/store.py). |
| Run: uv run pytest tests/test_audit_g5_data_store_store.py -m "not gpu" -v |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pytest |
|
|
| from align.data.store import ShardDataset, ShardReader, ShardWriter, read_shard_receipt |
|
|
|
|
| def _add(w, ids, pid, **extra): |
| w.add_row(ids, problem_id=pid, subject="Algebra", boundary_positions=[], **extra) |
|
|
|
|
| def _make_shard(d, name="s", rows=((1, 1, 1), (2, 2, 2, 2), (3, 3))): |
| with ShardWriter(d, name) as w: |
| for i, ids in enumerate(rows): |
| _add(w, list(ids), f"p{i}") |
| return w.receipt |
|
|
|
|
| |
|
|
|
|
| def test_u09_1_failed_add_row_leaves_no_partial_row(tmp_workdir): |
| """A row rejected by schema/JSON validation must not leave its tokens in |
| the blobs: the caller catches and continues (build.py's pattern), and the |
| NEXT row's metadata would be paired with the PREVIOUS row's tokens.""" |
| w = ShardWriter(tmp_workdir, "s") |
| _add(w, [1, 1, 1], "a", k=1) |
| with pytest.raises(ValueError): |
| _add(w, [2, 2, 2, 2], "b", k=1, extra_key=5) |
| with pytest.raises(TypeError): |
| _add(w, [9, 9, 9, 9, 9], "c", k=np.int64(3)) |
| _add(w, [3, 3], "d", k=2) |
| w.finalize() |
| rows = list(ShardReader(tmp_workdir, "s")[i] for i in range(2)) |
| assert [r.problem_id for r in rows] == ["a", "d"] |
| assert rows[0].input_ids.tolist() == [1, 1, 1] |
| assert rows[1].input_ids.tolist() == [3, 3] |
| assert w.receipt.n_tokens == 5 |
|
|
|
|
| def test_u09_1_non_json_extra_on_first_row_does_not_poison_schema(tmp_workdir): |
| w = ShardWriter(tmp_workdir, "s") |
| with pytest.raises(TypeError): |
| _add(w, [1, 1], "a", bad=np.int64(3)) |
| _add(w, [2, 2], "b", good=1) |
| w.finalize() |
| r = ShardReader(tmp_workdir, "s") |
| assert len(r) == 1 and r[0].problem_id == "b" and r[0].extras == {"good": 1} |
|
|
|
|
| def test_u09_1_ragged_field_failure_is_atomic_across_fields(tmp_workdir): |
| w = ShardWriter(tmp_workdir, "s", {"input_ids": "int32", "rej": "int32"}) |
| w.add_row( |
| [1, 1], problem_id="a", subject="x", boundary_positions=[], ragged={"rej": [5, 5]} |
| ) |
| with pytest.raises(ValueError): |
| w.add_row( |
| [2, 2], |
| problem_id="b", |
| subject="x", |
| boundary_positions=[], |
| ragged={"rej": [[1, 2], [3, 4]]}, |
| ) |
| w.add_row( |
| [3, 3, 3], problem_id="c", subject="x", boundary_positions=[], ragged={"rej": [6]} |
| ) |
| w.finalize() |
| r = ShardReader(tmp_workdir, "s") |
| assert [r[i].problem_id for i in range(len(r))] == ["a", "c"] |
| assert r[1].input_ids.tolist() == [3, 3, 3] and r[1].ragged["rej"].tolist() == [6] |
|
|
|
|
| |
|
|
|
|
| def test_u09_2_exception_in_with_block_leaves_no_receipt(tmp_workdir): |
| with pytest.raises(KeyboardInterrupt): |
| with ShardWriter(tmp_workdir, "s") as w: |
| _add(w, [1, 1, 1], "a") |
| raise KeyboardInterrupt |
| assert w.receipt is None |
| assert list(Path(tmp_workdir).glob("s.*")) == [] |
| with pytest.raises(FileNotFoundError): |
| ShardDataset(tmp_workdir) |
|
|
|
|
| def test_u09_2_clean_exit_still_finalizes(tmp_workdir): |
| r = _make_shard(tmp_workdir) |
| assert r.n_rows == 3 and (Path(tmp_workdir) / "s.receipt.json").exists() |
|
|
|
|
| def test_u09_2_finalized_shard_survives_a_later_exception(tmp_workdir): |
| with pytest.raises(RuntimeError): |
| with ShardWriter(tmp_workdir, "s") as w: |
| _add(w, [1, 1], "a") |
| w.finalize() |
| raise RuntimeError("after finalize") |
| assert (Path(tmp_workdir) / "s.receipt.json").exists() |
|
|
|
|
| def test_u09_2_poisoned_writer_refuses_to_finalize(tmp_workdir): |
| """A write that failed AFTER the first byte hit disk cannot be undone.""" |
| w = ShardWriter(tmp_workdir, "s") |
| _add(w, [1, 1], "a") |
| w._handles["input_ids"].close() |
| with pytest.raises(ValueError): |
| _add(w, [2, 2], "b") |
| with pytest.raises(RuntimeError, match="poison|abort|partial"): |
| _add(w, [3, 3], "c") |
| with pytest.raises(RuntimeError, match="poison|abort|partial"): |
| w.finalize() |
|
|
|
|
| |
|
|
|
|
| def test_u09_3_tampered_idx_is_detected(tmp_workdir): |
| _make_shard(tmp_workdir) |
| idx = Path(tmp_workdir) / "s.input_ids.idx.npy" |
| np.save(idx, np.array([0, 3, 7], dtype=np.int64)) |
| with pytest.raises((RuntimeError, ValueError)): |
| r = ShardReader(tmp_workdir, "s") |
| r[0] |
|
|
|
|
| def test_u09_3_trailing_bytes_in_bin_are_detected(tmp_workdir): |
| _make_shard(tmp_workdir) |
| with open(Path(tmp_workdir) / "s.input_ids.bin", "ab") as f: |
| f.write(b"\0" * 8) |
| with pytest.raises((RuntimeError, ValueError)): |
| r = ShardReader(tmp_workdir, "s") |
| r[0] |
|
|
|
|
| def test_u09_3_shifted_row_boundaries_change_the_hash(tmp_workdir): |
| """Same length, same last offset, different row boundaries: only a hash |
| that covers the idx files notices.""" |
| _make_shard(tmp_workdir) |
| r = ShardReader(tmp_workdir, "s") |
| r.verify() |
| idx = Path(tmp_workdir) / "s.input_ids.idx.npy" |
| np.save(idx, np.array([0, 2, 9, 9], dtype=np.int64)) |
| with pytest.raises(RuntimeError, match="hash"): |
| ShardReader(tmp_workdir, "s").verify() |
|
|
|
|
| def test_u09_3_legacy_v1_receipt_still_verifies(tmp_workdir): |
| """Shards built before the idx files were hashed carry format_version 1 and |
| must keep verifying against the file set they were hashed over.""" |
| from align.data.store import _hash_files |
|
|
| _make_shard(tmp_workdir) |
| d = Path(tmp_workdir) |
| legacy = _hash_files( |
| [d / f"s.{f}.bin" for f in ("boundary_positions", "input_ids")] + [d / "s.extras.jsonl"] |
| ) |
| rj = json.loads((d / "s.receipt.json").read_text()) |
| rj["content_hash"], rj["format_version"] = legacy, 1 |
| (d / "s.receipt.json").write_text(json.dumps(rj)) |
| ShardReader(tmp_workdir, "s").verify() |
|
|
|
|
| |
|
|
|
|
| def test_u09_4_rewriting_a_shard_drops_its_stale_receipt_immediately(tmp_workdir): |
| """Between `open(wb)` truncating the blobs and finalize(), the OLD receipt |
| must not be left describing bytes that no longer exist.""" |
| _make_shard(tmp_workdir) |
| w = ShardWriter(tmp_workdir, "s") |
| assert not (Path(tmp_workdir) / "s.receipt.json").exists() |
| w.abort() |
|
|
