diff --git a/.gitignore b/.gitignore index 21a9240..c31617e 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,6 @@ dist/ # IDE .idea/ .vscode/ + +# pytest +.pytest_cache/ diff --git a/pyproject.toml b/pyproject.toml index a3f33f4..4da1991 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..2dcbfaa --- /dev/null +++ b/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() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..95e198b --- /dev/null +++ b/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() diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..c3d1e0d --- /dev/null +++ b/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 diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..bb6bb38 --- /dev/null +++ b/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) diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..d648412 --- /dev/null +++ b/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() diff --git a/uv.lock b/uv.lock index 7721a1a..ba53679 100644 --- a/uv.lock +++ b/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"