refactor(tests): split into tests/unit and tests/tools

tests/tools/test_trainer.py

@@ -0,0 +1,286 @@
+"""Tests for Flow Matching training pipeline."""
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tools.trainer import load_and_prepare_data, train, _augment
+
+
+def _make_synthetic_trace(start, end, n_moves=30):
+    """Create a synthetic trace dict mimicking real JSONL format."""
+    sx, sy = start
+    ex, ey = end
+    dist = math.hypot(ex - sx, ey - sy)
+    angle = math.degrees(math.atan2(ey - sy, ex - sx))
+
+    events = []
+    for i in range(n_moves):
+        t_frac = i / (n_moves - 1)
+        x = int(sx + (ex - sx) * t_frac + np.random.normal(0, 3))
+        y = int(sy + (ey - sy) * t_frac + np.random.normal(0, 3))
+        t = int(t_frac * 500 + np.random.normal(0, 5))
+        events.append({"type": "move", "x": x, "y": y, "t": max(0, t)})
+
+    events.sort(key=lambda e: e["t"])
+    events[0]["t"] = 0
+
+    last_t = events[-1]["t"]
+    events.append({"type": "down", "x": ex, "y": ey, "t": last_t + 50})
+    events.append({"type": "up", "x": ex, "y": ey, "t": last_t + 130})
+
+    return {
+        "meta": {"start": [sx, sy], "end": [ex, ey], "dist": int(dist), "angle": round(angle, 1)},
+        "events": events,
+    }
+
+
+@pytest.fixture
+def synthetic_traces_file(tmp_path):
+    """Create a temp JSONL file with 25 synthetic traces."""
+    traces_path = tmp_path / "traces.jsonl"
+    rng = np.random.default_rng(42)
+    lines = []
+    for _ in range(25):
+        sx, sy = int(rng.integers(50, 750)), int(rng.integers(50, 750))
+        angle = rng.uniform(0, 2 * math.pi)
+        dist = int(rng.integers(100, 500))
+        ex = int(sx + dist * math.cos(angle))
+        ey = int(sy + dist * math.sin(angle))
+        ex = max(0, min(800, ex))
+        ey = max(0, min(600, ey))
+        trace = _make_synthetic_trace((sx, sy), (ex, ey))
+        lines.append(json.dumps(trace))
+    traces_path.write_text("\n".join(lines), encoding="utf-8")
+    return traces_path
+
+
+class TestLoadAndPrepare:
+    def test_returns_correct_shapes(self, synthetic_traces_file):
+        seq, cond, click_durs = load_and_prepare_data(synthetic_traces_file, seq_len=64)
+        assert seq.shape[1] == 64
+        assert seq.shape[2] == 3
+        assert cond.shape[1] == 3
+        assert len(seq) > 0
+
+    def test_forward_starts_near_zero(self, synthetic_traces_file):
+        seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
+        assert abs(seq[:, 0, 0].mean()) < 0.15
+
+    def test_forward_ends_near_one(self, synthetic_traces_file):
+        seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
+        assert abs(seq[:, -1, 0].mean() - 1.0) < 0.15
+
+
+class TestAugment:
+    def test_augmentation_multiplies_data(self, synthetic_traces_file):
+        seq, cond, _ = load_and_prepare_data(synthetic_traces_file, seq_len=64)
+        n_orig = len(seq)
+        seq_aug, cond_aug = _augment(seq, cond)
+        assert len(seq_aug) == n_orig * 6
+        assert len(cond_aug) == n_orig * 6
+
+
+class TestTrain:
+    def test_train_produces_model_files(self, synthetic_traces_file, tmp_path):
+        output_dir = tmp_path / "models"
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=output_dir,
+            epochs=3,
+            batch_size=8,
+            seq_len=64,
+        )
+        assert (output_dir / "flow_model.pt").exists()
+        assert (output_dir / "click_dist.json").exists()
+        assert (output_dir / "duration_dist.json").exists()
+        assert (output_dir / "train_config.json").exists()
+
+    def test_train_loss_decreases(self, synthetic_traces_file, tmp_path):
+        output_dir = tmp_path / "models"
+        losses = []
+
+        def cb(msg):
+            if "loss" in msg:
+                losses.append(msg["loss"])
+
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=output_dir,
+            epochs=20,
+            batch_size=8,
+            seq_len=64,
+            progress_callback=cb,
+        )
+        first_half = np.mean(losses[:10])
+        second_half = np.mean(losses[10:])
+        assert second_half < first_half
+
+
+class TestTrajectoryDataset:
+    def test_dataset_length_with_augmentation(self):
+        """Dataset length = N * 6 when augment=True."""
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((10, 64, 3), dtype=np.float32)
+        cond = np.zeros((10, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        assert len(ds) == 60
+
+    def test_dataset_length_without_augmentation(self):
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((10, 64, 3), dtype=np.float32)
+        cond = np.zeros((10, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=False)
+        assert len(ds) == 10
+
+    def test_getitem_returns_tensors(self):
+        from tools.trainer import TrajectoryDataset
+        import torch
+        seq = np.random.randn(5, 64, 3).astype(np.float32)
+        cond = np.random.randn(5, 3).astype(np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s, c = ds[0]
+        assert isinstance(s, torch.Tensor)
+        assert isinstance(c, torch.Tensor)
+        assert s.shape == (64, 3)
+        assert c.shape == (3,)
+
+    def test_aug_id_zero_returns_original(self):
+        """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
+        from tools.trainer import TrajectoryDataset
+        import torch
+        seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
+        cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s0, c0 = ds[0]
+        np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5)
+        np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5)
+
+    def test_aug_id_one_flips_lateral(self):
+        """Aug id 1 should flip the sign of the lateral channel (index 1)."""
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        seq[0, :, 1] = 0.5  # lateral all positive
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        # idx=1 → base_idx=0, aug_id=1 → flip
+        s1, _ = ds[1]
+        assert (s1[:, 1] < 0).all()
+
+    def test_aug_id_two_slows_speed(self):
+        """Aug id 2 should add log(1.25) to log_dt channel and cond[2]."""
+        import math
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s2, c2 = ds[2]  # idx=2 → base=0, aug_id=2
+        expected_delta = math.log(1.25)
+        np.testing.assert_allclose(s2[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c2[2].item() - expected_delta) < 1e-5
+
+    def test_aug_id_three_speeds_up(self):
+        """Aug id 3 should add log(1/1.2) to log_dt channel and cond[2]."""
+        import math
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s3, c3 = ds[3]  # idx=3 → base=0, aug_id=3
+        expected_delta = math.log(1.0 / 1.2)
+        np.testing.assert_allclose(s3[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c3[2].item() - expected_delta) < 1e-5
+
+    def test_aug_id_four_adds_temporal_noise(self):
+        """Aug id 4 should add Gaussian noise to log_dt (channel 2), leaving other channels unchanged."""
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s4, c4 = ds[4]  # idx=4 → base=0, aug_id=4
+        # Channels 0 and 1 should be unchanged (still 0)
+        np.testing.assert_allclose(s4[:, 0].numpy(), 0.0, atol=1e-6)
+        np.testing.assert_allclose(s4[:, 1].numpy(), 0.0, atol=1e-6)
+        # Channel 2 (log_dt) at position 0 is unchanged; positions 1+ have noise
+        assert s4[0, 2].item() == 0.0
+        # Noise should be non-zero (with overwhelming probability for 63 samples)
+        assert not np.allclose(s4[1:, 2].numpy(), 0.0)
+        # cond is unchanged for aug_id 4
+        np.testing.assert_allclose(c4.numpy(), cond[0], atol=1e-6)
+
+    def test_aug_id_five_flips_and_slows(self):
+        """Aug id 5 should flip lateral and add log(1/0.9) to log_dt and cond[2]."""
+        import math
+        from tools.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        seq[0, :, 1] = 0.5  # lateral positive
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s5, c5 = ds[5]  # idx=5 → base=0, aug_id=5
+        # Lateral flipped
+        assert (s5[:, 1] < 0).all()
+        # log_dt shifted
+        expected_delta = math.log(1.0 / 0.9)
+        np.testing.assert_allclose(s5[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c5[2].item() - expected_delta) < 1e-5
+
+
+class TestResumeFrom:
+    def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):
+        """train() with resume_from should load weights from given checkpoint dir."""
+        import torch
+        from tools.trainer import train
+        from tools.models import TrajectoryFlowModel
+
+        # First, train an initial model and save it
+        ckpt_dir = tmp_path / "pretrain"
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=ckpt_dir,
+            epochs=2,
+            batch_size=8,
+            seq_len=64,
+        )
+        assert (ckpt_dir / "flow_model.pt").exists()
+
+        # Read its weights to compare later
+        m_pretrain = TrajectoryFlowModel(seq_len=64)
+        m_pretrain.load_state_dict(torch.load(ckpt_dir / "flow_model.pt", weights_only=True))
+        first_param_pre = next(m_pretrain.parameters()).clone()
+
+        # Now train with resume_from for 1 epoch — weights should still be loaded
+        out_dir = tmp_path / "finetune"
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=out_dir,
+            epochs=1,
+            batch_size=8,
+            seq_len=64,
+            resume_from=ckpt_dir,
+        )
+
+        m_after = TrajectoryFlowModel(seq_len=64)
+        m_after.load_state_dict(torch.load(out_dir / "flow_model.pt", weights_only=True))
+        first_param_after = next(m_after.parameters())
+
+        # After 1 epoch, weights should be close to pre-train, not random init
+        # (random init would be O(1) magnitude apart; 1 epoch on small data shifts O(0.1))
+        diff = (first_param_pre - first_param_after).abs().mean().item()
+        assert diff < 0.5, f"Resume_from weights diverged too much: {diff}"
+
+    def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path):
+        from tools.trainer import train
+
+        with pytest.raises(FileNotFoundError):
+            train(
+                data_path=synthetic_traces_file,
+                output_dir=tmp_path / "out",
+                epochs=1,
+                batch_size=8,
+                seq_len=64,
+                resume_from=tmp_path / "nonexistent",
+            )