| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| import subprocess |
| import sys |
| from typing import cast |
| import pytest |
|
|
| from d2.split_audit import audit_plan, validate_manifest |
|
|
|
|
| def fixture() -> tuple[dict[str, object], list[dict[str, object]]]: |
| prompts = [{'pid': 'p1', 'split': 'train', 'family': 'holding_summary', |
| 'prompt_sha256': 'h1', 'source_docs': ['u1']}, |
| {'pid': 'p2', 'split': 'heldout_legal', 'family': 'chronology', |
| 'prompt_sha256': 'h2', 'source_docs': ['u2']}] |
| manifest: dict[str, object] = {'prompts': prompts, 'n_prompts': 2, |
| 'counts': {'train': {'holding_summary': 1}, 'heldout_legal': {'chronology': 1}}} |
| plan: list[dict[str, object]] = [{**p, 'seq_id': p['pid'], 'kind': 'gen'} for p in prompts] |
| return manifest, plan |
|
|
|
|
| def test_exact_population_and_partial_coverage() -> None: |
| m, plan = fixture() |
| r = audit_plan(plan[:1], m) |
| assert r['manifest_prompts'] == 2 and r['not_replayed_prompt_ids'] == ['p2'] |
| assert r['included_prompt_ids'] == ['p1'] |
|
|
|
|
| @pytest.mark.parametrize('key,value', [('seq_id', 'unknown'), ('split', 'dev'), |
| ('family', 'argument_section'), ('prompt_sha256', 'other')]) |
| def test_refuses_foreign_exact_identity(key: str, value: str) -> None: |
| m, plan = fixture() |
| plan[0][key] = value |
| with pytest.raises(ValueError, match='differs'): |
| audit_plan(plan, m) |
|
|
|
|
| def test_refuses_duplicate_plan_and_acap_overlap() -> None: |
| m, plan = fixture() |
| with pytest.raises(ValueError, match='duplicate'): |
| audit_plan(plan + plan[:1], m) |
| plan[0].update(kind='acap', split='acap', family='acap', tokens_sha256='t1') |
| with pytest.raises(ValueError, match='acap'): |
| audit_plan(plan[:1], m) |
|
|
|
|
| def test_two_captures_of_one_acap_prompt(tmp_path: Path) -> None: |
| m, _ = fixture() |
| captures = tmp_path / 'acap' |
| captures.mkdir() |
| |
| for i in range(2): |
| capture = {'schema': 'd2.acap.v1', 'seq_id': f'acap{i}', 'prompt_sha256': 'capture-prompt', |
| 'gen_start': 20, 'tokens': list(range(20)) + list(range(30 + i, 40 + i))} |
| (captures / f'acap{i}.json').write_text(json.dumps(capture)) |
| out = tmp_path / 'plan' |
| subprocess.run([sys.executable, '-m', 'd2.replay', 'plan', '--seqs', '', |
| '--acap', str(captures), '--out', str(out)], check=True, capture_output=True) |
| plan = cast(list[dict[str, object]], [json.loads(line) for line in (out / 'plan.jsonl').read_text().splitlines()]) |
| assert plan[0]['prompt_sha256'] == plan[1]['prompt_sha256'] |
| assert plan[0]['tokens_sha256'] != plan[1]['tokens_sha256'] |
| assert audit_plan(plan, m)['by_split'] == {'acap': 2} |
| plan[1]['tokens_sha256'] = plan[0]['tokens_sha256'] |
| with pytest.raises(ValueError, match='duplicate'): |
| audit_plan(plan, m) |
| plan[1]['tokens_sha256'] = 'different-capture' |
| plan[1]['prompt_sha256'] = 'h2' |
| with pytest.raises(ValueError, match='acap'): |
| audit_plan(plan, m) |
|
|
|
|
| def test_refuses_source_family_leak_and_stale_counts() -> None: |
| m, _ = fixture() |
| prompts = m['prompts'] |
| assert isinstance(prompts, list) |
| prompts[1]['source_docs'] = ['u1', 'u2'] |
| with pytest.raises(ValueError, match='source family'): |
| validate_manifest(m) |
| m, _ = fixture() |
| m['n_prompts'] = 3 |
| with pytest.raises(ValueError, match='n_prompts'): |
| validate_manifest(m) |
| m, _ = fixture() |
| m['counts'] = {} |
| with pytest.raises(ValueError, match='counts'): |
| validate_manifest(m) |
|
|
