101 lines
4.1 KiB
Python
101 lines
4.1 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from tools 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
|