| import json |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from dlb.adapters.base import AdapterError |
| from dlb.adapters.mdlm import MDLMAdapter |
|
|
|
|
| def _capture_record(index: int) -> dict[str, object]: |
| prefix = [1] * 64 |
| continuation = [2] * 64 |
| return { |
| "sample_id": index, |
| "text": "fixed prefix generated continuation", |
| "token_ids": prefix + continuation, |
| "prompt_id": index, |
| "completion_id": 0, |
| "source_index": index + 100, |
| "prefix_token_ids": prefix, |
| "reference_token_ids": [3] * 64, |
| "full_token_ids": prefix + continuation, |
| "prefix_text": "fixed prefix", |
| "continuation_text": "generated continuation", |
| "reference_text": "source continuation", |
| "full_text": "fixed prefix generated continuation", |
| "prefix_exact_match": True, |
| } |
|
|
|
|
| def _write_capture(path: Path, samples: list[dict[str, object]]) -> None: |
| path.write_text( |
| json.dumps({"schema": "dlb-upstream-token-capture-v1", "samples": samples}), |
| encoding="utf-8", |
| ) |
|
|
|
|
| def test_conditional_capture_reader_accepts_complete_prompt_records(tmp_path: Path) -> None: |
| """Catch converter regressions that drop prompt/source metadata from captures.""" |
|
|
| path = tmp_path / "upstream_token_ids.json" |
| _write_capture(path, [_capture_record(0), _capture_record(1)]) |
|
|
| records = MDLMAdapter()._read_conditional_capture(path, 2) |
|
|
| assert [record["prompt_id"] for record in records] == [0, 1] |
| assert records[0]["prefix_exact_match"] is True |
|
|
|
|
| @pytest.mark.parametrize( |
| ("mutate", "message"), |
| [ |
| (lambda record: record.update(prefix_exact_match=False), "prefix mismatch"), |
| (lambda record: record.update(extra="bad"), "record format"), |
| (lambda record: record.update(sample_id=1), "expected sample_id 0"), |
| ], |
| ) |
| def test_conditional_capture_reader_rejects_invalid_prompt_records( |
| tmp_path: Path, mutate, message: str |
| ) -> None: |
| """Catch malformed conditional captures before publication JSONL is written.""" |
|
|
| path = tmp_path / "upstream_token_ids.json" |
| record = _capture_record(0) |
| mutate(record) |
| _write_capture(path, [record]) |
|
|
| with pytest.raises(AdapterError, match=message): |
| MDLMAdapter()._read_conditional_capture(path, 1) |
|
|
