Add pytest suite with 40 unit and integration tests

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.

.gitignore

@@ -17,3 +17,6 @@ dist/
 # IDE
 .idea/
 .vscode/
+
+# pytest
+.pytest_cache/

pyproject.toml

@@ -30,6 +30,13 @@ mouse-collect = "mouse_control.collect:main"
 mouse-train = "mouse_control.train:main"
 mouse-visualize = "mouse_control.visualize:main"
 
+[dependency-groups]
+dev = ["pytest>=8.0"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/mouse_control"]
 artifacts = ["src/mouse_control/assets/*.onnx"]

tests/conftest.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = REPO_ROOT / "data"
+
+
+@pytest.fixture(scope="session")
+def train_csv() -> Path:
+    path = DATA_DIR / "mouse_data.csv"
+    if not path.is_file():
+        pytest.skip(f"missing training csv: {path}")
+    return path
+
+
+@pytest.fixture(scope="session")
+def test_csv() -> Path:
+    path = DATA_DIR / "mouse_data_test.csv"
+    if not path.is_file():
+        pytest.skip(f"missing test csv: {path}")
+    return path
+
+
+@pytest.fixture(scope="session")
+def bundled_session():
+    """Shared inference session for the bundled model (loading is slow)."""
+    from mouse_control import load_model
+
+    return load_model()

tests/test_cli.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+CLI_COMMANDS = ["mouse-collect", "mouse-train", "mouse-visualize"]
+
+
+@pytest.mark.parametrize("command", CLI_COMMANDS)
+def test_cli_help(command: str) -> None:
+    """Every CLI must respond to --help with exit 0 and an argparse usage line."""
+    result = subprocess.run(
+        [command, "--help"], capture_output=True, text=True, check=False,
+    )
+    assert result.returncode == 0, result.stderr
+    assert "usage:" in result.stdout
+
+
+def test_visualize_no_show_produces_png(tmp_path: Path) -> None:
+    save_path = tmp_path / "out.png"
+    result = subprocess.run(
+        [
+            "mouse-visualize", "--no-show", "--seed", "1",
+            "--n-trajectories", "3", "--n-interp", "20",
+            "--save", str(save_path),
+        ],
+        capture_output=True, text=True, check=False,
+    )
+    assert result.returncode == 0, result.stderr
+    assert save_path.exists()
+    assert save_path.stat().st_size > 1000
+
+
+def test_visualize_module_invocation(tmp_path: Path) -> None:
+    """The module is also runnable via `python -m mouse_control.visualize`."""
+    save_path = tmp_path / "out.png"
+    result = subprocess.run(
+        [
+            sys.executable, "-m", "mouse_control.visualize",
+            "--no-show", "--seed", "2", "--n-trajectories", "2",
+            "--save", str(save_path),
+        ],
+        capture_output=True, text=True, check=False,
+    )
+    assert result.returncode == 0, result.stderr
+    assert save_path.exists()

tests/test_inference.py

@@ -0,0 +1,95 @@
+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

tests/test_model.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from mouse_control import SimpleNet
+
+
+@pytest.fixture
+def model() -> SimpleNet:
+    return SimpleNet()
+
+
+@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
+def test_forward_with_unsqueezed_input(model: SimpleNet, batch_size: int) -> None:
+    out = model(torch.randn(batch_size, 1, 2))
+    assert out.shape == (batch_size, 10, 2)
+
+
+@pytest.mark.parametrize("batch_size", [1, 8, 32])
+def test_forward_with_flat_input(model: SimpleNet, batch_size: int) -> None:
+    """SimpleNet's first layer flattens, so (N, 2) is also valid."""
+    out = model(torch.randn(batch_size, 2))
+    assert out.shape == (batch_size, 10, 2)
+
+
+def test_layers_have_expected_dimensions(model: SimpleNet) -> None:
+    assert model.fc1.in_features == 2
+    assert model.fc1.out_features == 64
+    assert model.fc2.in_features == 64
+    assert model.fc2.out_features == 32
+    assert model.fc3.in_features == 32
+    assert model.fc3.out_features == 20
+
+
+def test_model_has_trainable_parameters(model: SimpleNet) -> None:
+    params = list(model.parameters())
+    assert params, "SimpleNet should expose parameters"
+    assert all(p.requires_grad for p in params)
+
+
+def test_forward_is_differentiable(model: SimpleNet) -> None:
+    inp = torch.randn(2, 1, 2, requires_grad=True)
+    out = model(inp)
+    out.sum().backward()
+    assert inp.grad is not None
+    assert inp.grad.shape == inp.shape
+
+
+def test_onnx_export_round_trips(tmp_path, model: SimpleNet) -> None:
+    """Trained or not, the model architecture must round-trip via ONNX."""
+    import numpy as np
+    import onnxruntime as ort
+
+    onnx_path = tmp_path / "round_trip.onnx"
+    dummy = torch.randn(1, 1, 2)
+    model.eval()
+    torch.onnx.export(
+        model, dummy, str(onnx_path),
+        input_names=["input"], output_names=["output"],
+        external_data=False,
+    )
+    assert onnx_path.exists()
+    assert not onnx_path.with_suffix(onnx_path.suffix + ".data").exists(), \
+        "external_data=False must produce a single-file ONNX"
+
+    session = ort.InferenceSession(str(onnx_path))
+    inp = np.array([[[42.0, -17.0]]], dtype=np.float32)
+    out = session.run(None, {"input": inp})[0]
+    assert out.shape == (1, 10, 2)
+
+    with torch.no_grad():
+        torch_out = model(torch.from_numpy(inp)).numpy()
+    np.testing.assert_allclose(out, torch_out, atol=1e-4)

tests/test_train.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+
+from mouse_control.train import TrajectoryDataset, _load_csv, train
+
+
+def _write_csv(path: Path, rows: list[list[tuple[int, int]]]) -> None:
+    """Write rows in the project's quirky single-quoted-pair format."""
+    with open(path, "w", newline="") as f:
+        writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
+        for row in rows:
+            # First column is the target (last keypoint), then 10 keypoints.
+            target_x, target_y = row[-1]
+            cells = [f"{target_x},{target_y}"] + [f"{x},{y}" for x, y in row]
+            writer.writerow(cells)
+
+
+@pytest.fixture
+def synthetic_csv(tmp_path: Path) -> Path:
+    """Tiny linear trajectories from origin -> target."""
+    path = tmp_path / "synth.csv"
+    rows = []
+    rng = np.random.default_rng(0)
+    for _ in range(20):
+        tx, ty = int(rng.integers(-150, 151)), int(rng.integers(-150, 151))
+        keypoints = [
+            (int(round(tx * i / 9)), int(round(ty * i / 9))) for i in range(10)
+        ]
+        rows.append(keypoints)
+    _write_csv(path, rows)
+    return path
+
+
+def test_load_csv_shapes(synthetic_csv: Path) -> None:
+    inputs, labels = _load_csv(synthetic_csv)
+    assert inputs.shape == (20, 1, 2)
+    assert labels.shape == (20, 10, 2)
+    assert inputs.dtype == torch.float32
+    assert labels.dtype == torch.float32
+
+
+def test_load_csv_round_trip(synthetic_csv: Path) -> None:
+    """The target column (input) must equal the last keypoint of the labels."""
+    inputs, labels = _load_csv(synthetic_csv)
+    last_keypoint = labels[:, -1, :]
+    np.testing.assert_array_equal(inputs.squeeze(1).numpy(), last_keypoint.numpy())
+
+
+def test_dataset_indexing(synthetic_csv: Path) -> None:
+    inputs, labels = _load_csv(synthetic_csv)
+    ds = TrajectoryDataset(inputs, labels)
+    assert len(ds) == 20
+    x, y = ds[0]
+    assert x.shape == (1, 2)
+    assert y.shape == (10, 2)
+
+
+def test_load_csv_missing_file_raises(tmp_path: Path) -> None:
+    with pytest.raises(FileNotFoundError):
+        _load_csv(tmp_path / "missing.csv")
+
+
+def test_train_end_to_end_produces_single_file_onnx(
+    synthetic_csv: Path, tmp_path: Path,
+) -> None:
+    """A quick training run must export a valid, self-contained ONNX model."""
+    output = tmp_path / "trained.onnx"
+    train(
+        train_csv=synthetic_csv,
+        test_csv=synthetic_csv,
+        output=output,
+        epochs=3,
+        batch_size=8,
+    )
+    assert output.exists()
+    assert output.stat().st_size > 1000
+    # external_data=False -> no .data sidecar
+    assert not output.with_suffix(output.suffix + ".data").exists()
+
+    session = ort.InferenceSession(str(output))
+    inp = np.array([[[100.0, 50.0]]], dtype=np.float32)
+    out = session.run(None, {"input": inp})[0]
+    assert out.shape == (1, 10, 2)
+
+
+def test_train_with_real_data(train_csv: Path, test_csv: Path, tmp_path: Path) -> None:
+    """Smoke test against the project's actual data."""
+    output = tmp_path / "real.onnx"
+    train(
+        train_csv=train_csv,
+        test_csv=test_csv,
+        output=output,
+        epochs=2,
+        batch_size=32,
+    )
+    assert output.exists()

uv.lock

@@ -22,6 +22,15 @@ resolution-markers = [
     "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'",
 ]
 
+[[package]]
+name = "colorama"
+version = "0.4.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
+]
+
 [[package]]
 name = "contourpy"
 version = "1.3.3"
@@ -233,6 +242,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d5/0c/043d5e551459da400957a1395e0febbf771446ff34291afcbe3d8be2a279/fsspec-2026.4.0-py3-none-any.whl", hash = "sha256:11ef7bb35dab8a394fde6e608221d5cf3e8499401c249bebaeaad760a1a8dec2", size = 203402, upload-time = "2026-04-29T20:42:36.842Z" },
 ]
 
+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
+]
+
 [[package]]
 name = "jinja2"
 version = "3.1.6"
@@ -498,6 +516,11 @@ dependencies = [
     { name = "torch" },
 ]
 
+[package.dev-dependencies]
+dev = [
+    { name = "pytest" },
+]
+
 [package.metadata]
 requires-dist = [
     { name = "matplotlib", specifier = ">=3.10" },
@@ -509,6 +532,9 @@ requires-dist = [
     { name = "torch", specifier = ">=2.5" },
 ]
 
+[package.metadata.requires-dev]
+dev = [{ name = "pytest", specifier = ">=8.0" }]
+
 [[package]]
 name = "mpmath"
 version = "1.3.0"
@@ -961,6 +987,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" },
 ]
 
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
+]
+
 [[package]]
 name = "protobuf"
 version = "7.34.1"
@@ -976,6 +1011,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/95/608f665226bca68b736b79e457fded9a2a38c4f4379a4a7614303d9db3bc/protobuf-7.34.1-py3-none-any.whl", hash = "sha256:bb3812cd53aefea2b028ef42bd780f5b96407247f20c6ef7c679807e9d188f11", size = 170715, upload-time = "2026-03-20T17:34:45.384Z" },
 ]
 
+[[package]]
+name = "pygments"
+version = "2.20.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" },
+]
+
 [[package]]
 name = "pyparsing"
 version = "3.3.2"
@@ -985,6 +1029,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" },
 ]
 
+[[package]]
+name = "pytest"
+version = "9.0.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/549bd94f1a0a402dc8cf64563a117c0f3765662e2e668477624baeec44d5/pytest-9.0.3.tar.gz", hash = "sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c", size = 1572165, upload-time = "2026-04-07T17:16:18.027Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d4/24/a372aaf5c9b7208e7112038812994107bc65a84cd00e0354a88c2c77a617/pytest-9.0.3-py3-none-any.whl", hash = "sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9", size = 375249, upload-time = "2026-04-07T17:16:16.13Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"