116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from tools.runtime import runtime_assets as assets
|
|
|
|
|
|
@pytest.fixture
|
|
def bundle(tmp_path):
|
|
files = {
|
|
"model/config.json": b"{}",
|
|
"model/modules.json": b"[]",
|
|
"model/tokenizer.json": b"{}",
|
|
"model/model.safetensors": b"weights",
|
|
"tiktoken/token-file": b"tokenizer",
|
|
}
|
|
for name, value in files.items():
|
|
path = tmp_path / name
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_bytes(value)
|
|
manifest = {
|
|
"schema_version": 1,
|
|
"embedding": {"model": assets.DEFAULT_MODEL, "revision": assets.DEFAULT_REVISION},
|
|
"encodings": ["cl100k_base"],
|
|
"files": {name: hashlib.sha256(value).hexdigest() for name, value in files.items()},
|
|
}
|
|
(tmp_path / "manifest.json").write_text(json.dumps(manifest))
|
|
return tmp_path, manifest
|
|
|
|
|
|
def test_assets_verify_integrity_and_reject_corruption(bundle):
|
|
root, manifest = bundle
|
|
assert assets.verify(root) == manifest
|
|
(root / "model/model.safetensors").write_bytes(b"corrupted")
|
|
with pytest.raises(RuntimeError, match="poškodený"):
|
|
assets.verify(root)
|
|
|
|
|
|
def test_prepared_assets_override_empty_user_cache_and_pin_revision(bundle, monkeypatch, tmp_path):
|
|
root, _ = bundle
|
|
monkeypatch.setenv("ZP_RUNTIME_ASSETS", str(root))
|
|
monkeypatch.setenv("ZP_ASSETS_REQUIRED", "1")
|
|
monkeypatch.setenv("TIKTOKEN_CACHE_DIR", str(tmp_path / "empty"))
|
|
assets.configure_tokenizer()
|
|
assert os.environ["TIKTOKEN_CACHE_DIR"] == str(root / "tiktoken")
|
|
assert assets.local_model(assets.DEFAULT_MODEL, None) == str(root / "model")
|
|
with pytest.raises(RuntimeError, match="revízia"):
|
|
assets.local_model(assets.DEFAULT_MODEL, "a" * 40)
|
|
with pytest.raises(RuntimeError, match="nie je v image"):
|
|
assets.local_model("other/model", None)
|
|
|
|
|
|
def test_required_assets_never_silently_fall_back_to_remote_download(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("ZP_RUNTIME_ASSETS", str(tmp_path))
|
|
monkeypatch.setenv("ZP_ASSETS_REQUIRED", "1")
|
|
with pytest.raises(RuntimeError, match="Chýbajú runtime assets"):
|
|
assets.configure_tokenizer()
|
|
|
|
|
|
def test_existing_assets_are_verified_without_download(bundle, monkeypatch):
|
|
root, manifest = bundle
|
|
probes = []
|
|
monkeypatch.setattr(assets, "offline_probe", lambda path: probes.append(path))
|
|
assert assets.prepare(root, assets.DEFAULT_MODEL, None, ["cl100k_base"]) == manifest
|
|
assert probes == [root]
|
|
|
|
|
|
def test_manifest_cannot_escape_asset_directory(bundle):
|
|
root, manifest = bundle
|
|
manifest["files"]["../outside"] = "a" * 64
|
|
(root / "manifest.json").write_text(json.dumps(manifest))
|
|
with pytest.raises(RuntimeError, match="mimo runtime assets"):
|
|
assets.verify(root)
|
|
|
|
|
|
def test_fresh_api_lifespan_loads_model_and_tokenizer_without_network_or_user_cache(tmp_path):
|
|
environment = {
|
|
**os.environ,
|
|
"ZP_RUNTIME_ASSETS": str(assets.asset_root().resolve()),
|
|
"ZP_ASSETS_REQUIRED": "1",
|
|
"HF_HOME": str(tmp_path / "empty-hf"),
|
|
"TIKTOKEN_CACHE_DIR": str(tmp_path / "empty-tokenizer"),
|
|
"HF_HUB_OFFLINE": "1",
|
|
"TRANSFORMERS_OFFLINE": "1",
|
|
"TEST_JOB_DIR": str(tmp_path / "jobs"),
|
|
}
|
|
code = """
|
|
import socket, asyncio, os
|
|
from pathlib import Path
|
|
def denied(*args, **kwargs): raise AssertionError('Unexpected network access')
|
|
socket.socket.connect = denied
|
|
socket.socket.connect_ex = denied
|
|
socket.getaddrinfo = denied
|
|
import app.main as main
|
|
from tools.rag.build_chunks import TOKEN_COUNTER
|
|
from tools.rag.embedding_utils import embed_query
|
|
main.DATA_DIR = Path(os.environ['TEST_JOB_DIR'])
|
|
async def run():
|
|
async with main.lifespan(main.app):
|
|
assert main.app.state.embedding_ready
|
|
assert TOKEN_COUNTER.encoding is not None
|
|
assert embed_query('štart bez siete').shape[0] == 384
|
|
asyncio.run(run())
|
|
print('offline API lifespan OK')
|
|
"""
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", code], env=environment, capture_output=True, text=True, timeout=60
|
|
)
|
|
assert result.returncode == 0, result.stderr
|
|
assert "offline API lifespan OK" in result.stdout
|