Huang Qi
3c1cf171a7
Provide bundled_path() and resolve() helpers that locate ONNX weights and JSON metadata via importlib.resources, falling back to a user-supplied directory. Missing assets raise ModelLoadError. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
35 lines
936 B
Python
35 lines
936 B
Python
"""Test the asset path resolver."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from ai_mouse import _assets
|
|
from ai_mouse.errors import ModelLoadError
|
|
|
|
|
|
def test_bundled_flow_model_exists() -> None:
|
|
p = _assets.bundled_path("flow_model.onnx")
|
|
assert p.exists()
|
|
assert p.suffix == ".onnx"
|
|
|
|
|
|
def test_bundled_train_config_loadable() -> None:
|
|
p = _assets.bundled_path("train_config.json")
|
|
cfg = json.loads(p.read_text())
|
|
assert "seq_len" in cfg
|
|
assert "d_model" in cfg
|
|
|
|
|
|
def test_resolve_with_custom_dir(tmp_path: Path) -> None:
|
|
(tmp_path / "flow_model.onnx").write_bytes(b"x")
|
|
p = _assets.resolve(tmp_path, "flow_model.onnx")
|
|
assert p == tmp_path / "flow_model.onnx"
|
|
|
|
|
|
def test_missing_asset_raises_model_load_error(tmp_path: Path) -> None:
|
|
with pytest.raises(ModelLoadError, match="missing"):
|
|
_assets.resolve(tmp_path, "nonexistent.onnx")
|