48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
"""Tests for model registry (src/model_registry.py)."""
|
|
|
|
import pytest
|
|
from src.model_registry import scan_models, ModelConfig
|
|
|
|
|
|
def test_scan_returns_list():
|
|
result = scan_models("/nonexistent/path")
|
|
assert isinstance(result, list)
|
|
|
|
|
|
def test_scan_empty_dir(tmp_path):
|
|
result = scan_models(str(tmp_path))
|
|
assert result == []
|
|
|
|
|
|
def test_scan_finds_pt_files(tmp_path):
|
|
(tmp_path / "best.pt").write_bytes(b"fake")
|
|
(tmp_path / "truck-detector.pt").write_bytes(b"fake")
|
|
result = scan_models(str(tmp_path))
|
|
assert len(result) == 2
|
|
names = {m.filename for m in result}
|
|
assert "best.pt" in names
|
|
assert "truck-detector.pt" in names
|
|
|
|
|
|
def test_scan_skips_non_model_files(tmp_path):
|
|
(tmp_path / "modelREADME.md").write_text("readme")
|
|
(tmp_path / "best.pt").write_bytes(b"fake")
|
|
result = scan_models(str(tmp_path))
|
|
assert len(result) == 1
|
|
|
|
|
|
def test_model_config_fields(tmp_path):
|
|
(tmp_path / "v4-best.pt").write_bytes(b"fake")
|
|
result = scan_models(str(tmp_path))
|
|
cfg = result[0]
|
|
assert cfg.filename == "v4-best.pt"
|
|
assert cfg.path == str(tmp_path / "v4-best.pt")
|
|
assert isinstance(cfg.known_classes, list)
|
|
|
|
|
|
def test_model_config_fallback_classes(tmp_path):
|
|
(tmp_path / "unknown-model.pt").write_bytes(b"fake")
|
|
result = scan_models(str(tmp_path))
|
|
cfg = result[0]
|
|
assert cfg.known_classes == []
|