File size: 4,660 Bytes
f340984 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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)
|