tests/test_model_store.py

import hashlib

import pytest

import unlimited_ocr_rdna4.model_store as model_store
from unlimited_ocr_rdna4.errors import ModelIntegrityError


def _synthetic_upstream() -> str:
    lines = [
        "from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM",
        "cor_list = eval(ref_text[2])",
        "lines = eval(outputs)['Line']['line']",
        "line_type = eval(outputs)['Line']['line_type']",
        "endpoints = eval(outputs)['Line']['line_endpoint']",
        "p0 = eval(line.split(' -- ')[0])",
        "p1 = eval(line.split(' -- ')[-1])",
        "(x, y) = eval(endpoint.split(': ')[1])",
        "images_seq_mask[idx].unsqueeze(-1).cuda()",
    ]
    for _ in range(3):
        lines.extend(
            [
                "                input_ids=input_ids.unsqueeze(0).cuda(),",
                "                eos_token_id=tokenizer.eos_token_id,",
            ]
        )
    return "\n".join(lines) + "\n"


def test_patch_is_exact_and_idempotent(monkeypatch) -> None:
    source = _synthetic_upstream()
    monkeypatch.setattr(model_store, "UPSTREAM_MODEL_CODE_SHA256", model_store.sha256_text(source))
    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", "__TO_BE_FILLED__")
    patched = model_store.patch_model_source_text(source)
    patched_hash = model_store.sha256_text(patched)
    assert patched.count("ast.literal_eval(") == 7
    assert patched.count("attention_mask=torch.ones_like") == 3
    assert patched.count("pad_token_id=tokenizer.eos_token_id") == 3
    assert ".to(inputs_embeds.device)" in patched

    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", patched_hash)
    assert model_store.patch_model_source_text(patched) == patched


def test_patch_rejects_unknown_source() -> None:
    with pytest.raises(ModelIntegrityError, match="refusing to patch unknown model code"):
        model_store.patch_model_source_text("unknown")


def test_verify_model_reports_missing_directory(tmp_path) -> None:
    with pytest.raises(ModelIntegrityError, match="real directory"):
        model_store.verify_model(tmp_path / "missing")


def test_verify_model_with_small_fixture(tmp_path, monkeypatch) -> None:
    code = "patched model code\n"
    weight = b"weights"
    (tmp_path / "modeling_unlimitedocr.py").write_text(code, encoding="utf-8")
    (tmp_path / "weight.bin").write_bytes(weight)
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(weight))
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_SHA256", hashlib.sha256(weight).hexdigest())
    monkeypatch.setattr(model_store, "PATCHED_MODEL_CODE_SHA256", hashlib.sha256(code.encode()).hexdigest())
    monkeypatch.setattr(
        model_store,
        "MODEL_PAYLOAD_SHA256",
        {
            "modeling_unlimitedocr.py": hashlib.sha256(code.encode()).hexdigest(),
            "weight.bin": hashlib.sha256(weight).hexdigest(),
        },
    )
    model_store._write_manifest(tmp_path)
    status = model_store.verify_model(tmp_path)
    assert status.prepared
    assert status.weight_sha256 == hashlib.sha256(weight).hexdigest()


def test_verify_model_rejects_unexpected_file(tmp_path, monkeypatch) -> None:
    code = b"code"
    weight = b"weight"
    (tmp_path / "code.py").write_bytes(code)
    (tmp_path / "weight.bin").write_bytes(weight)
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(weight))
    monkeypatch.setattr(
        model_store,
        "MODEL_PAYLOAD_SHA256",
        {"code.py": hashlib.sha256(code).hexdigest(), "weight.bin": hashlib.sha256(weight).hexdigest()},
    )
    model_store._write_manifest(tmp_path)
    (tmp_path / "configuration_surprise.py").write_text("raise SystemExit", encoding="utf-8")
    with pytest.raises(ModelIntegrityError, match="unexpected"):
        model_store.verify_model(tmp_path)


def test_verify_model_rejects_symlinked_payload(tmp_path, monkeypatch) -> None:
    target = tmp_path / "target"
    target.write_bytes(b"weight")
    (tmp_path / "weight.bin").symlink_to(target)
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_FILE", "weight.bin")
    monkeypatch.setattr(model_store, "MODEL_WEIGHT_BYTES", len(b"weight"))
    monkeypatch.setattr(
        model_store,
        "MODEL_PAYLOAD_SHA256",
        {"weight.bin": hashlib.sha256(b"weight").hexdigest(), "target": hashlib.sha256(b"weight").hexdigest()},
    )
    model_store._write_manifest(tmp_path)
    with pytest.raises(ModelIntegrityError, match="missing regular files"):
        model_store.verify_model(tmp_path)