Huang Qi
c8fa5db7d3
Coverage: - test_model: SimpleNet forward (parametrized over batch sizes and both unsqueezed and flat input shapes), layer dimensions, differentiability, and ONNX round-trip - test_inference: load_model resolution order (bundled, cwd override, explicit path, missing path), and predict shape/dtype/determinism plus endpoint sanity across 8 cardinal/diagonal targets - test_train: _load_csv parsing, TrajectoryDataset indexing, full train() pipeline producing a single-file ONNX, plus a smoke test against the real data shipped under data/ - test_cli: --help for the three console scripts and a real run of mouse-visualize via both the entry point and python -m Wire up pytest via dependency-groups and tool.pytest.ini_options.
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import shutil
|
|
from importlib.resources import as_file, files
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
import pytest
|
|
|
|
from mouse_control import load_model, predict
|
|
|
|
|
|
def _bundled_path() -> Path:
|
|
"""Resolve the bundled mouse.onnx to a stable filesystem path for tests."""
|
|
with as_file(files("mouse_control.assets").joinpath("mouse.onnx")) as p:
|
|
return Path(p)
|
|
|
|
|
|
def test_load_model_returns_inference_session(bundled_session) -> None:
|
|
assert isinstance(bundled_session, ort.InferenceSession)
|
|
|
|
|
|
def test_bundled_input_output_names(bundled_session) -> None:
|
|
assert bundled_session.get_inputs()[0].name == "input"
|
|
assert bundled_session.get_outputs()[0].name == "output"
|
|
|
|
|
|
def test_load_model_explicit_path() -> None:
|
|
session = load_model(_bundled_path())
|
|
assert isinstance(session, ort.InferenceSession)
|
|
|
|
|
|
def test_load_model_missing_path_raises() -> None:
|
|
with pytest.raises(Exception):
|
|
load_model("/no/such/path.onnx")
|
|
|
|
|
|
def test_load_model_cwd_override(tmp_path, monkeypatch) -> None:
|
|
"""When ./mouse.onnx exists, it should take precedence over the bundle."""
|
|
monkeypatch.chdir(tmp_path)
|
|
shutil.copy(_bundled_path(), tmp_path / "mouse.onnx")
|
|
session = load_model() # implicit
|
|
assert isinstance(session, ort.InferenceSession)
|
|
# Must produce the same output as the bundle since the file is identical.
|
|
pkg_session = load_model(_bundled_path())
|
|
a = predict(session, 75, -40)
|
|
b = predict(pkg_session, 75, -40)
|
|
np.testing.assert_array_equal(a, b)
|
|
|
|
|
|
def test_load_model_falls_back_to_bundled(tmp_path, monkeypatch) -> None:
|
|
"""With no ./mouse.onnx, load_model() must succeed via the bundle."""
|
|
monkeypatch.chdir(tmp_path)
|
|
assert not (tmp_path / "mouse.onnx").exists()
|
|
session = load_model()
|
|
assert isinstance(session, ort.InferenceSession)
|
|
|
|
|
|
def test_predict_shape_and_dtype(bundled_session) -> None:
|
|
out = predict(bundled_session, 100.0, 200.0)
|
|
assert out.shape == (10, 2)
|
|
assert out.dtype == np.float32
|
|
|
|
|
|
def test_predict_is_deterministic(bundled_session) -> None:
|
|
a = predict(bundled_session, 100, 100)
|
|
b = predict(bundled_session, 100, 100)
|
|
np.testing.assert_array_equal(a, b)
|
|
|
|
|
|
def test_predict_int_and_float_args_equivalent(bundled_session) -> None:
|
|
a = predict(bundled_session, 50, 50)
|
|
b = predict(bundled_session, 50.0, 50.0)
|
|
np.testing.assert_array_equal(a, b)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dx,dy",
|
|
[
|
|
(150, 0), (-150, 0), (0, 150), (0, -150),
|
|
(100, 100), (-100, 100), (100, -100), (-100, -100),
|
|
],
|
|
)
|
|
def test_predict_endpoint_near_target(bundled_session, dx: int, dy: int) -> None:
|
|
"""Trained model's endpoint should fall within ~50px of the target."""
|
|
end = predict(bundled_session, dx, dy)[-1]
|
|
err = float(np.linalg.norm(end - np.array([dx, dy])))
|
|
assert err < 50, f"target ({dx},{dy}) endpoint err={err:.1f}"
|
|
|
|
|
|
def test_predict_zero_displacement_stays_near_origin(bundled_session) -> None:
|
|
end = predict(bundled_session, 0, 0)[-1]
|
|
assert abs(end[0]) < 20
|
|
assert abs(end[1]) < 20
|