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.

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()