File size: 2,293 Bytes
dc72977 | 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 | 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)
|