dougvk/Unlimited-OCR-RDNA4

1import hashlib2 3import pytest4 5import unlimited_ocr_rdna4.model_store as model_store6from unlimited_ocr_rdna4.errors import ModelIntegrityError7 8 9def _synthetic_upstream() -> str:10    lines = [11        "from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM",12        "cor_list = eval(ref_text[2])",13        "lines = eval(outputs)['Line']['line']",14        "line_type = eval(outputs)['Line']['line_type']",15        "endpoints = eval(outputs)['Line']['line_endpoint']",16        "p0 = eval(line.split(' -- ')[0])",17        "p1 = eval(line.split(' -- ')[-1])",18        "(x, y) = eval(endpoint.split(': ')[1])",19        "images_seq_mask[idx].unsqueeze(-1).cuda()",20    ]21    for _ in range(3):22        lines.extend(23            [24                "                input_ids=input_ids.unsqueeze(0).cuda(),",25                "                eos_token_id=tokenizer.eos_token_id,",26            ]27        )28    return "\n".join(lines) + "\n"29 30 31def test_patch_is_exact_and_idempotent(monkeypatch) -> None:32    source = _synthetic_upstream()33    monkeypatch.setattr(model_store, "UPSTREAM_MODEL_CODE_SHA256", model_store.sha256_text(source))34    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", "__TO_BE_FILLED__")35    patched = model_store.patch_model_source_text(source)36    patched_hash = model_store.sha256_text(patched)37    assert patched.count("ast.literal_eval(") == 738    assert patched.count("attention_mask=torch.ones_like") == 339    assert patched.count("pad_token_id=tokenizer.eos_token_id") == 340    assert ".to(inputs_embeds.device)" in patched41 42    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", patched_hash)43    assert model_store.patch_model_source_text(patched) == patched44 45 46def test_patch_rejects_unknown_source() -> None:47    with pytest.raises(ModelIntegrityError, match="refusing to patch unknown model code"):48        model_store.patch_model_source_text("unknown")49 50 51def test_verify_model_reports_missing_directory(tmp_path) -> None:52    with pytest.raises(ModelIntegrityError, match="real directory"):53        model_store.verify_model(tmp_path / "missing")54 55 56def test_verify_model_with_small_fixture(tmp_path, monkeypatch) -> None:57    code = "patched model code\n"58    weight = b"weights"59    (tmp_path / "modeling_unlimitedocr.py").write_text(code, encoding="utf-8")60    (tmp_path / "weight.bin").write_bytes(weight)61    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")62    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(weight))63    monkeypatch.setattr(model_store, "MODEL_WEIGHT_SHA256", hashlib.sha256(weight).hexdigest())64    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", hashlib.sha256(code.encode()).hexdigest())65    monkeypatch.setattr(66        model_store,67        "MODEL_PAYLOAD_SHA256",68        {69            "modeling_unlimitedocr.py": hashlib.sha256(code.encode()).hexdigest(),70            "weight.bin": hashlib.sha256(weight).hexdigest(),71        },72    )73    model_store._write_manifest(tmp_path)74    status = model_store.verify_model(tmp_path)75    assert status.prepared76    assert status.weight_sha256 == hashlib.sha256(weight).hexdigest()77 78 79def test_verify_model_rejects_unexpected_file(tmp_path, monkeypatch) -> None:80    code = b"code"81    weight = b"weight"82    (tmp_path / "code.py").write_bytes(code)83    (tmp_path / "weight.bin").write_bytes(weight)84    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")85    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(weight))86    monkeypatch.setattr(87        model_store,88        "MODEL_PAYLOAD_SHA256",89        {"code.py": hashlib.sha256(code).hexdigest(), "weight.bin": hashlib.sha256(weight).hexdigest()},90    )91    model_store._write_manifest(tmp_path)92    (tmp_path / "configuration_surprise.py").write_text("raise SystemExit", encoding="utf-8")93    with pytest.raises(ModelIntegrityError, match="unexpected"):94        model_store.verify_model(tmp_path)95 96 97def test_verify_model_rejects_symlinked_payload(tmp_path, monkeypatch) -> None:98    target = tmp_path / "target"99    target.write_bytes(b"weight")100    (tmp_path / "weight.bin").symlink_to(target)101    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")102    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(b"weight"))103    monkeypatch.setattr(104        model_store,105        "MODEL_PAYLOAD_SHA256",106        {"weight.bin": hashlib.sha256(b"weight").hexdigest(), "target": hashlib.sha256(b"weight").hexdigest()},107    )108    model_store._write_manifest(tmp_path)109    with pytest.raises(ModelIntegrityError, match="missing regular files"):110        model_store.verify_model(tmp_path)111