yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Huang Qi
2026-05-12 00:47:00 +08:00
parent b0f0d6db11
commit fc7b2bd236
15 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/tools/__init__.py Normal file
View File

56
tests/tools/conftest.py Normal file
View File

@@ -0,0 +1,56 @@
"""Shared test fixtures for ai_mouse."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pytest
import torch
from tools.models import TrajectoryFlowModel
from tools.scroll.models import ScrollCVAE
@pytest.fixture
def model_dir(tmp_path: Path) -> Path:
"""Create a temporary directory with trained Flow model artifacts."""
# Flow model
model = TrajectoryFlowModel(seq_len=64)
torch.save(model.state_dict(), tmp_path / "flow_model.pt")
# Click distribution
click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
(tmp_path / "click_dist.json").write_text(json.dumps(click_dist))
# Duration distribution
dur_dist = {
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
"params": [{"mu_log": 5.5, "sigma_log": 0.5}] * 8,
}
(tmp_path / "duration_dist.json").write_text(json.dumps(dur_dist))
# Train config (architecture params)
train_cfg = {
"seq_len": 64,
"d_model": 128,
"nhead": 4,
"num_layers": 4,
"dim_feedforward": 256,
"cond_dim": 3,
}
(tmp_path / "train_config.json").write_text(json.dumps(train_cfg))
return tmp_path
@pytest.fixture
def scroll_model_dir(tmp_path: Path) -> Path:
"""Create a temporary directory with trained scroll model artifacts."""
model = ScrollCVAE(seq_len=32)
torch.save(model.state_dict(), tmp_path / "scroll_model.pt")
scroll_cfg = {"seq_len": 32, "latent_dim": 16, "hidden": 64, "cond_dim": 7}
(tmp_path / "scroll_config.json").write_text(json.dumps(scroll_cfg))
return tmp_path

336
tests/tools/test_balabit_adapter.py Normal file
View File

@@ -0,0 +1,336 @@
"""Tests for Balabit Mouse Dynamics Challenge data adapter."""
from __future__ import annotations
import json
from pathlib import Path
import pytest
def test_module_exports():
"""The adapter module must export the public functions used by CLI."""
from tools.data_adapters import balabit
assert hasattr(balabit, "parse_session_csv")
assert hasattr(balabit, "segment_by_clicks")
assert hasattr(balabit, "filter_segments")
assert hasattr(balabit, "process_session")
assert hasattr(balabit, "MouseEvent")
assert hasattr(balabit, "Segment")
def test_mouse_event_dataclass():
"""MouseEvent has expected fields."""
from tools.data_adapters.balabit import MouseEvent
e = MouseEvent(t_ms=100, button="NoButton", state="Move", x=300, y=400)
assert e.t_ms == 100
assert e.state == "Move"
assert e.x == 300
def test_segment_dataclass():
"""Segment has expected fields."""
from tools.data_adapters.balabit import MouseEvent, Segment
events = [MouseEvent(t_ms=0, button="NoButton", state="Move", x=10, y=20)]
s = Segment(events=events, click_x=100, click_y=200, click_t_ms=500, session_id="user1_s1")
assert s.events == events
assert s.click_x == 100
assert s.session_id == "user1_s1"
def _write_csv(path: Path, rows: list[str]) -> None:
"""Helper to write a Balabit-format CSV with header."""
header = "record timestamp,client timestamp,button,state,x,y"
path.write_text(header + "\n" + "\n".join(rows) + "\n", encoding="utf-8")
class TestParseSessionCsv:
def test_parses_basic_rows(self, tmp_path):
from tools.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_1"
_write_csv(csv, [
"1500000000.000,0.000,NoButton,Move,100,200",
"1500000000.050,0.050,NoButton,Move,110,210",
"1500000000.100,0.100,Left,Pressed,120,220",
])
events = parse_session_csv(csv)
assert len(events) == 3
assert events[0].t_ms == 0
assert events[0].state == "Move"
assert events[0].x == 100
assert events[2].t_ms == 100
assert events[2].state == "Pressed"
assert events[2].button == "Left"
def test_handles_float_timestamps(self, tmp_path):
"""Client timestamps are floats in seconds; we convert to int ms."""
from tools.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_2"
_write_csv(csv, [
"0,1.234,NoButton,Move,50,60",
"0,1.250,NoButton,Move,55,65",
])
events = parse_session_csv(csv)
assert events[0].t_ms == 1234
assert events[1].t_ms == 1250
def test_skips_malformed_rows(self, tmp_path):
"""Rows with bad data are logged and skipped, not raised."""
from tools.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_3"
_write_csv(csv, [
"0,0.000,NoButton,Move,100,200",
"BROKEN_ROW",
"0,abc,NoButton,Move,100,200", # bad timestamp
"0,0.100,NoButton,Move,150,250",
])
events = parse_session_csv(csv)
assert len(events) == 2
assert events[0].x == 100
assert events[1].x == 150
def test_returns_empty_list_for_empty_file(self, tmp_path):
from tools.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_4"
csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8")
events = parse_session_csv(csv)
assert events == []
class TestSegmentByClicks:
def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"):
from tools.data_adapters.balabit import MouseEvent
return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y)
def test_one_click_one_segment(self):
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Move", 50, 60),
self._make_event(500, "Move", 100, 100),
self._make_event(600, "Pressed", 110, 110, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="test_s1")
assert len(segments) == 1
seg = segments[0]
assert seg.click_x == 110
assert seg.click_y == 110
assert seg.click_t_ms == 600
assert len(seg.events) == 3
assert seg.session_id == "test_s1"
def test_window_excludes_old_events(self):
"""Move events earlier than (click_t - window_ms) are dropped."""
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20), # too old
self._make_event(100, "Move", 20, 30), # too old
self._make_event(900, "Move", 30, 40), # in window
self._make_event(1000, "Pressed", 40, 50, button="Left"),
]
segments = segment_by_clicks(events, window_ms=200, session_id="s")
assert len(segments) == 1
assert len(segments[0].events) == 1
assert segments[0].events[0].t_ms == 900
def test_multiple_clicks_multiple_segments(self):
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Pressed", 50, 50, button="Left"),
self._make_event(200, "Released", 50, 50, button="Left"),
self._make_event(300, "Move", 60, 60),
self._make_event(400, "Move", 70, 70),
self._make_event(500, "Pressed", 80, 80, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 2
assert segments[0].click_x == 50
assert segments[1].click_x == 80
# Second segment's events must not include the first Pressed
for e in segments[1].events:
assert e.state == "Move"
def test_skips_pressed_with_non_left_button(self):
"""Right-clicks and wheel-clicks don't anchor segments (only Left)."""
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Pressed", 50, 50, button="Right"), # ignored
self._make_event(200, "Move", 60, 60),
self._make_event(300, "Pressed", 70, 70, button="Left"), # anchor
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 1
assert segments[0].click_x == 70
def test_no_clicks_returns_empty(self):
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Move", 20, 30),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert segments == []
def test_excludes_drag_events(self):
"""Drag events are not Move; segment should only include Move."""
from tools.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Drag", 30, 40), # not Move
self._make_event(200, "Move", 50, 60),
self._make_event(300, "Pressed", 70, 80, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 1
# Drag event should not appear in seg.events
assert all(e.state == "Move" for e in segments[0].events)
assert len(segments[0].events) == 2
class TestFilterSegments:
def _seg(self, events_data: list[tuple[int, int, int]], click=(500, 500), session="s"):
"""events_data: list of (t_ms, x, y) tuples."""
from tools.data_adapters.balabit import MouseEvent, Segment
events = [MouseEvent(t_ms=t, button="NoButton", state="Move", x=x, y=y)
for (t, x, y) in events_data]
click_t = events[-1].t_ms + 50 if events else 100
return Segment(events=events, click_x=click[0], click_y=click[1],
click_t_ms=click_t, session_id=session)
def test_drops_segment_with_too_few_events(self):
from tools.data_adapters.balabit import filter_segments
seg = self._seg([(0, 100, 100), (50, 105, 105)]) # 2 events, min=5
result = filter_segments([seg], min_events=5, min_dist=10,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_short_distance(self):
from tools.data_adapters.balabit import filter_segments
# Start (100,100), end click=(105,100) → dist=5, min_dist=50
events = [(i*10, 100+i, 100) for i in range(10)]
seg = self._seg(events, click=(105, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_too_long_span(self):
from tools.data_adapters.balabit import filter_segments
# Span 6000ms, max_span=5000
events = [(0, 100, 100), (2000, 200, 200), (4000, 300, 300), (6000, 400, 400),
(6010, 410, 400), (6020, 420, 400)]
seg = self._seg(events, click=(500, 500))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=10000)
assert result == []
def test_drops_segment_with_gap(self):
from tools.data_adapters.balabit import filter_segments
# Gap of 500ms between events 2 and 3, max_gap=200
events = [(0, 100, 100), (50, 110, 110), (100, 120, 120),
(600, 200, 200), (650, 210, 210), (700, 220, 220)]
seg = self._seg(events, click=(300, 300))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_out_of_range_coords(self):
from tools.data_adapters.balabit import filter_segments
# x=10000 out of range
events = [(i*10, 10000, 100) for i in range(6)]
seg = self._seg(events, click=(10100, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_short_arc_length(self):
"""Total arc < 50 even though endpoints are far → high-frequency jitter only."""
from tools.data_adapters.balabit import filter_segments
# Tiny back-and-forth with click 100px away — unrealistic, drop it
events = [(i*10, 100 + (i % 2), 100) for i in range(10)]
seg = self._seg(events, click=(200, 100))
# arc length = sum of |Δp| ≈ 9 (alternating ±1) which is < 50
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_keeps_valid_segment(self):
from tools.data_adapters.balabit import filter_segments
# Smooth 100→500 px straight line, 10 events, span 500ms
events = [(i*50, 100 + i*40, 100) for i in range(10)]
seg = self._seg(events, click=(500, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert len(result) == 1
class TestProcessSession:
def test_writes_jsonl_in_expected_format(self, tmp_path):
from tools.config import BalabitAdapterConfig
from tools.data_adapters.balabit import process_session
# Construct a Balabit-format CSV with one valid segment
csv_path = tmp_path / "user_session_42"
rows = []
# 10 Move events going from (100,100) to (500,200) over 500ms
for i in range(10):
t = i * 0.05
x = 100 + i * 40
y = 100 + i * 10
rows.append(f"0,{t:.3f},NoButton,Move,{x},{y}")
rows.append(f"0,0.550,Left,Pressed,510,210")
_write_csv(csv_path, rows)
out = tmp_path / "out.jsonl"
config = BalabitAdapterConfig(
window_ms=1200, min_dist=50, min_events=5,
max_span_ms=5000, max_gap_ms=200,
)
n = process_session(csv_path, out, config)
assert n == 1
assert out.exists()
line = out.read_text(encoding="utf-8").strip()
record = json.loads(line)
assert record["meta"]["start"] == [100, 100]
assert record["meta"]["end"] == [510, 210]
assert record["meta"]["dist"] > 0
assert record["meta"]["source"] == "balabit"
assert record["meta"]["session_id"] == "user_session_42"
# Events: only move events, no down/up
assert all(e["type"] == "move" for e in record["events"])
# Timestamps relative (start at 0)
assert record["events"][0]["t"] == 0
def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
from tools.config import BalabitAdapterConfig
from tools.data_adapters.balabit import process_session
csv_path = tmp_path / "empty_session"
_write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"]) # no clicks
out = tmp_path / "out.jsonl"
config = BalabitAdapterConfig()
n = process_session(csv_path, out, config)
assert n == 0
# File may not exist or may be empty — both acceptable
if out.exists():
assert out.read_text() == ""
def test_appends_to_existing_jsonl(self, tmp_path):
from tools.config import BalabitAdapterConfig
from tools.data_adapters.balabit import process_session
out = tmp_path / "out.jsonl"
out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")
csv_path = tmp_path / "session_x"
rows = [f"0,{i*0.05:.3f},NoButton,Move,{100+i*40},100" for i in range(10)]
rows.append("0,0.550,Left,Pressed,510,100")
_write_csv(csv_path, rows)
config = BalabitAdapterConfig()
process_session(csv_path, out, config)
lines = out.read_text(encoding="utf-8").strip().split("\n")
assert len(lines) == 2 # original + appended

155
tests/tools/test_eval_metrics.py Normal file
View File

@@ -0,0 +1,155 @@
"""Tests for the eval metrics module."""
from __future__ import annotations
import numpy as np
import pytest
class TestKinematics:
def test_compute_speed_constant_velocity(self):
"""Constant-velocity trajectory has constant speed."""
from tools.eval.metrics import compute_speed
# 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms
xs = np.arange(0, 100, 10, dtype=float)
ys = np.zeros(10, dtype=float)
ts = np.arange(0, 1000, 100, dtype=float)
v = compute_speed(xs, ys, ts)
# All speeds should be ≈ 0.1 px/ms
assert v.shape == (9,) # n-1 differences
np.testing.assert_allclose(v, 0.1, rtol=1e-4)
def test_compute_speed_handles_zero_dt(self):
"""Adjacent points with same timestamp must not produce NaN/inf."""
from tools.eval.metrics import compute_speed
xs = np.array([0.0, 10.0, 20.0])
ys = np.array([0.0, 0.0, 0.0])
ts = np.array([0.0, 0.0, 100.0]) # zero dt between [0] and [1]
v = compute_speed(xs, ys, ts)
assert np.isfinite(v).all()
def test_compute_acceleration(self):
"""Linearly increasing speed → constant acceleration."""
from tools.eval.metrics import compute_acceleration
# speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms²
speeds = np.array([0.1, 0.2, 0.3, 0.4])
ts = np.array([100.0, 200.0, 300.0, 400.0])
a = compute_acceleration(speeds, ts)
np.testing.assert_allclose(a, 0.001, rtol=1e-4)
def test_compute_jerk(self):
from tools.eval.metrics import compute_jerk
# accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001
accels = np.array([0.001, 0.002, 0.003])
ts = np.array([200.0, 300.0, 400.0])
j = compute_jerk(accels, ts)
np.testing.assert_allclose(j, 1e-5, rtol=1e-4)
class TestStatsSummary:
def test_compute_stats_returns_expected_keys(self):
from tools.eval.metrics import compute_stats
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
s = compute_stats(x)
assert "mean" in s
assert "std" in s
assert "cv" in s
assert "p25" in s
assert "p50" in s
assert "p75" in s
assert "p95" in s
def test_cv_for_constant_is_zero(self):
from tools.eval.metrics import compute_stats
x = np.full(10, 3.0)
s = compute_stats(x)
assert s["cv"] == 0.0
class TestFftSpectrum:
def test_finds_dominant_frequency(self):
"""A pure 8 Hz signal should have its peak near 8 Hz."""
from tools.eval.metrics import fft_spectrum
# Sample at 100 Hz for 1 second
sample_rate_hz = 100.0
ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz)
signal = np.sin(2 * np.pi * 8 * ts_ms / 1000) # 8 Hz sine
freqs, mags = fft_spectrum(signal, sample_rate_hz)
peak_freq = freqs[np.argmax(mags)]
assert abs(peak_freq - 8.0) < 1.0 # within 1 Hz
def test_returns_only_positive_frequencies(self):
from tools.eval.metrics import fft_spectrum
signal = np.random.randn(64)
freqs, mags = fft_spectrum(signal, 50.0)
assert (freqs >= 0).all()
assert len(freqs) == len(mags)
class TestKlDivergence:
def test_identical_distributions_zero_kl(self):
"""KL(p, p) ≈ 0."""
from tools.eval.metrics import kl_divergence_histograms
rng = np.random.default_rng(42)
x = rng.normal(0, 1, 5000)
y = rng.normal(0, 1, 5000)
kl = kl_divergence_histograms(x, y, bins=50)
assert kl < 0.05
def test_different_distributions_positive_kl(self):
"""Different means → positive KL."""
from tools.eval.metrics import kl_divergence_histograms
rng = np.random.default_rng(42)
x = rng.normal(0, 1, 5000)
y = rng.normal(3, 1, 5000)
kl = kl_divergence_histograms(x, y, bins=50)
assert kl > 0.5
def test_handles_disjoint_supports(self):
"""No NaN even when histograms have non-overlapping bins."""
from tools.eval.metrics import kl_divergence_histograms
x = np.array([1.0, 1.1, 1.2, 1.3, 1.4])
y = np.array([10.0, 10.1, 10.2, 10.3, 10.4])
kl = kl_divergence_histograms(x, y, bins=10)
assert np.isfinite(kl)
class TestReportGeneration:
def test_generates_report_md(self, tmp_path):
"""Smoke test: build_report writes an MD file with all expected sections."""
from tools.eval.report import build_report
# Synthetic generated traces (3 traces, 50 points each)
rng = np.random.default_rng(0)
gen_traces = []
for _ in range(3):
xs = np.cumsum(rng.uniform(0, 5, 50))
ys = np.cumsum(rng.uniform(-1, 1, 50))
ts = np.cumsum(rng.uniform(5, 20, 50))
gen_traces.append({"xs": xs, "ys": ys, "ts": ts})
# Synthetic reference
ref_traces = []
for _ in range(5):
xs = np.cumsum(rng.uniform(0, 5, 50))
ys = np.cumsum(rng.uniform(-1, 1, 50))
ts = np.cumsum(rng.uniform(5, 20, 50))
ref_traces.append({"xs": xs, "ys": ys, "ts": ts})
out_md = tmp_path / "report.md"
build_report(
generated_traces=gen_traces,
reference_traces=ref_traces,
output_md=out_md,
tag="smoke-test",
model_dir="/fake/model/dir",
)
assert out_md.exists()
content = out_md.read_text(encoding="utf-8")
assert "# Eval Report" in content
assert "smoke-test" in content
assert "速度" in content or "speed" in content.lower()
assert "FFT" in content.upper()
# PNG plots should exist next to MD
plot_dir = tmp_path / "plots"
assert plot_dir.exists()
assert any(plot_dir.iterdir())

69
tests/tools/test_models.py Normal file
View File

@@ -0,0 +1,69 @@
"""Tests for TrajectoryFlowModel architecture."""
from __future__ import annotations
import torch
import pytest
from tools.models import TrajectoryFlowModel
class TestTrajectoryFlowModel:
"""Test the Conditional Flow Matching model."""
@pytest.fixture
def model(self):
return TrajectoryFlowModel(
seq_len=64, d_model=128, nhead=4, num_layers=4,
dim_feedforward=256, dropout=0.1, cond_dim=3,
)
def test_output_shape(self, model):
"""(4, 64, 3) input → (4, 64, 3) output."""
batch = 4
x_t = torch.randn(batch, 64, 3)
t = torch.rand(batch)
cond = torch.randn(batch, 3)
out = model(x_t, t, cond)
assert out.shape == (batch, 64, 3)
def test_single_sample(self, model):
"""(1, 64, 3) works."""
x_t = torch.randn(1, 64, 3)
t = torch.rand(1)
cond = torch.randn(1, 3)
out = model(x_t, t, cond)
assert out.shape == (1, 64, 3)
def test_deterministic(self, model):
"""Eval mode, same input → same output."""
model.eval()
x_t = torch.randn(2, 64, 3)
t = torch.tensor([0.3, 0.7])
cond = torch.randn(2, 3)
with torch.no_grad():
out1 = model(x_t, t, cond)
out2 = model(x_t, t, cond)
torch.testing.assert_close(out1, out2)
def test_different_timesteps(self, model):
"""t=0.1 vs t=0.9 gives different output."""
model.eval()
x_t = torch.randn(1, 64, 3)
cond = torch.randn(1, 3)
with torch.no_grad():
out_early = model(x_t, torch.tensor([0.1]), cond)
out_late = model(x_t, torch.tensor([0.9]), cond)
assert not torch.allclose(out_early, out_late, atol=1e-5)
def test_gradient_flows(self, model):
"""Backward works, grad on x_t exists."""
model.train()
x_t = torch.randn(2, 64, 3, requires_grad=True)
t = torch.rand(2)
cond = torch.randn(2, 3)
out = model(x_t, t, cond)
loss = out.sum()
loss.backward()
assert x_t.grad is not None
assert x_t.grad.shape == (2, 64, 3)
assert x_t.grad.abs().sum() > 0

65
tests/tools/test_scroll_collector.py Normal file
View File

@@ -0,0 +1,65 @@
"""Tests for scroll collection state and target generation."""
from __future__ import annotations
import pytest
from tools.scroll.collector import ScrollCollector
class TestNextTarget:
def test_target_mode_distance_range(self):
sc = ScrollCollector(mode="target", count=10, page_height=10000, viewport_height=900)
for _ in range(20):
result = sc.next_target(current_scrollY=2000)
dist = abs(result["target_scrollY"] - 2000)
assert 500 <= dist <= 3000
assert result["direction"] in ("up", "down")
def test_fast_mode_distance_range(self):
sc = ScrollCollector(mode="fast", count=10, page_height=10000, viewport_height=900)
for _ in range(20):
result = sc.next_target(current_scrollY=5000)
dist = abs(result["target_scrollY"] - 5000)
assert 3000 <= dist <= 8000
def test_precise_mode_distance_range(self):
sc = ScrollCollector(mode="precise", count=10, page_height=10000, viewport_height=900)
for _ in range(20):
result = sc.next_target(current_scrollY=3000)
dist = abs(result["target_scrollY"] - 3000)
assert 200 <= dist <= 800
def test_target_within_page_bounds(self):
sc = ScrollCollector(mode="target", count=10, page_height=10000, viewport_height=900)
result = sc.next_target(current_scrollY=1000)
assert 0 <= result["target_scrollY"] <= 10000
result = sc.next_target(current_scrollY=9000)
assert 0 <= result["target_scrollY"] <= 10000
def test_success_zone_by_mode(self):
sc = ScrollCollector(mode="target", count=10, page_height=10000)
assert sc.success_radius == 80
sc2 = ScrollCollector(mode="fast", count=10, page_height=10000)
assert sc2.success_radius == 120
sc3 = ScrollCollector(mode="precise", count=10, page_height=10000)
assert sc3.success_radius == 40
def test_target_always_reachable(self):
"""Target must always be reachable: user can scroll to bring it into success zone."""
sc = ScrollCollector(mode="target", count=10, page_height=10000, viewport_height=900)
viewport_center = 450 # 900 / 2
max_scroll_top = 10000 - 900 # 9100
for current in [0, 100, 500, 2000, 5000, 8000, 9000]:
for _ in range(10):
result = sc.next_target(current_scrollY=current)
target = result["target_scrollY"]
# The scrollTop needed to bring target into viewport center
needed_scroll = target - viewport_center + 25
# Must be achievable (0 <= needed_scroll <= max_scroll_top)
# With success_radius=80, there's a window, not exact match needed
reachable_min = viewport_center - sc.success_radius
reachable_max = max_scroll_top + viewport_center + sc.success_radius
assert reachable_min <= target <= reachable_max, (
f"Target {target} not reachable from scrollY={current}"
)

37
tests/tools/test_scroll_models.py Normal file
View File

@@ -0,0 +1,37 @@
"""Tests for ScrollCVAE model."""
from __future__ import annotations
import torch
import pytest
from tools.scroll.models import ScrollCVAE
class TestScrollCVAEForward:
@pytest.fixture
def model(self):
return ScrollCVAE(seq_len=32, latent_dim=16, hidden=64, cond_dim=7)
def test_output_shapes(self, model):
batch = 4
seq = torch.randn(batch, 32, 2)
cond = torch.randn(batch, 7)
recon, mu, logvar = model(seq, cond)
assert recon.shape == (batch, 32, 2)
assert mu.shape == (batch, 16)
assert logvar.shape == (batch, 16)
def test_decode_shape(self, model):
z = torch.randn(4, 16)
cond = torch.randn(4, 7)
out = model.decode(z, cond)
assert out.shape == (4, 32, 2)
def test_decode_deterministic(self, model):
model.eval()
z = torch.randn(1, 16)
cond = torch.randn(1, 7)
with torch.no_grad():
out1 = model.decode(z, cond)
out2 = model.decode(z, cond)
torch.testing.assert_close(out1, out2)

86
tests/tools/test_scroll_trainer.py Normal file
View File

@@ -0,0 +1,86 @@
"""Tests for scroll training pipeline."""
from __future__ import annotations
import json
import math
from pathlib import Path
import numpy as np
import pytest
from tools.scroll.trainer import load_scroll_data, train_scroll, _augment_scroll
def _make_synthetic_scroll_trace(mode="target"):
"""Create a synthetic scroll trace."""
distance = {"target": 1500, "fast": 5000, "precise": 400}[mode]
direction = "down"
start = 2000
target = start + distance
events = []
n_events = 20
for i in range(n_events):
frac = (i + 1) / n_events
delta = int(distance / n_events * (1 + 0.2 * np.random.randn()))
delta = max(20, delta)
t = int(frac * 800 + np.random.normal(0, 10))
events.append({"deltaY": delta, "deltaMode": 0, "t": max(0, t)})
events.sort(key=lambda e: e["t"])
events[0]["t"] = 0
return {
"meta": {
"mode": mode,
"start_scrollY": start,
"target_scrollY": target,
"end_scrollY": target + 5,
"distance": distance,
"direction": direction,
"duration_ms": events[-1]["t"],
"viewport_height": 900,
},
"events": events,
}
@pytest.fixture
def synthetic_scroll_file(tmp_path):
traces_path = tmp_path / "scroll_traces.jsonl"
lines = []
for mode in ["target", "fast", "precise"]:
for _ in range(10):
lines.append(json.dumps(_make_synthetic_scroll_trace(mode)))
traces_path.write_text("\n".join(lines), encoding="utf-8")
return traces_path
class TestLoadScrollData:
def test_returns_correct_shapes(self, synthetic_scroll_file):
seq, cond = load_scroll_data(synthetic_scroll_file, seq_len=32)
assert seq.shape[1] == 32
assert seq.shape[2] == 2 # (delta_norm, log_dt)
assert cond.shape[1] == 7
assert len(seq) > 0
class TestAugment:
def test_4x_augmentation(self, synthetic_scroll_file):
seq, cond = load_scroll_data(synthetic_scroll_file, seq_len=32)
n = len(seq)
seq_aug, cond_aug = _augment_scroll(seq, cond)
assert len(seq_aug) == n * 4
class TestTrainScroll:
def test_produces_model_files(self, synthetic_scroll_file, tmp_path):
output_dir = tmp_path / "scroll_models"
train_scroll(
data_path=synthetic_scroll_file,
output_dir=output_dir,
epochs=3,
batch_size=8,
)
assert (output_dir / "scroll_model.pt").exists()
assert (output_dir / "scroll_config.json").exists()

180
tests/tools/test_server.py Normal file
View File

@@ -0,0 +1,180 @@
"""Integration tests for the ai_mouse server API routes."""
from __future__ import annotations
import json
import pytest
import pytest_asyncio
from httpx import ASGITransport, AsyncClient
from tools.server import create_app
from tools.server.deps import get_data_dir
@pytest.fixture
def app():
return create_app()
@pytest_asyncio.fixture
async def client(app):
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as c:
yield c
# ---------------------------------------------------------------------------
# Status endpoint
# ---------------------------------------------------------------------------
class TestStatus:
@pytest.mark.asyncio
async def test_status_returns_trace_count(self, client):
resp = await client.get("/api/status")
assert resp.status_code == 200
data = resp.json()
assert "trace_count" in data
assert "model_trained" in data
assert isinstance(data["trace_count"], int)
assert isinstance(data["model_trained"], bool)
# ---------------------------------------------------------------------------
# Collect endpoints
# ---------------------------------------------------------------------------
class TestCollect:
@pytest.mark.asyncio
async def test_start_returns_ab_positions(self, client):
resp = await client.post(
"/api/collect/start",
json={"count": 10, "dist_min": 50, "dist_max": 400},
)
assert resp.status_code == 200
data = resp.json()
assert "a" in data
assert "b" in data
assert len(data["a"]) == 2
assert len(data["b"]) == 2
@pytest.mark.asyncio
async def test_skip_returns_new_positions(self, client):
# Start first
await client.post(
"/api/collect/start",
json={"count": 10, "dist_min": 50, "dist_max": 400},
)
resp = await client.post("/api/collect/skip")
assert resp.status_code == 200
data = resp.json()
assert "a" in data
assert "b" in data
@pytest.mark.asyncio
async def test_trace_without_start_returns_400(self, client):
# Reset state by creating a fresh app
resp = await client.post(
"/api/collect/trace",
json={"meta": {}, "events": []},
)
# May or may not be 400 depending on state from other tests
# Just verify the endpoint is reachable
assert resp.status_code in (200, 400)
@pytest.mark.asyncio
async def test_collect_trace_increments_count(self, client, tmp_path, monkeypatch):
"""Test that posting a trace increments the collected count."""
# Monkeypatch data dir to use tmp
import tools.server.deps as deps
monkeypatch.setattr(deps, "_DATA_DIR", tmp_path)
# Start collection
await client.post(
"/api/collect/start",
json={"count": 5, "dist_min": 50, "dist_max": 400},
)
# Post a trace
trace = {
"meta": {"start": [100, 200], "end": [300, 400], "dist": 283, "angle": 45},
"events": [
{"type": "move", "x": 100, "y": 200, "t": 0},
{"type": "move", "x": 200, "y": 300, "t": 50},
{"type": "down", "x": 300, "y": 400, "t": 100},
{"type": "up", "x": 300, "y": 400, "t": 180},
],
}
resp = await client.post("/api/collect/trace", json=trace)
assert resp.status_code == 200
data = resp.json()
assert data["collected"] == 1
assert data["remaining"] == 4
assert data["a"] is not None
assert data["b"] is not None
# ---------------------------------------------------------------------------
# Verify endpoint
# ---------------------------------------------------------------------------
class TestVerify:
@pytest.mark.asyncio
async def test_verify_returns_paths(self, client, model_dir, monkeypatch):
"""Test trajectory generation endpoint."""
import tools.server.routes_verify as rv
# We can't easily monkeypatch the model dir used inside the route
# but we can test the endpoint is accessible
resp = await client.post(
"/api/verify",
json={"start": [100, 100], "end": [500, 300], "n_paths": 2},
)
# Will fail with 404 if no models exist - that's expected in test env
# We just verify the endpoint routes correctly
assert resp.status_code in (200, 404, 500)
# ---------------------------------------------------------------------------
# Scroll endpoints
# ---------------------------------------------------------------------------
class TestScroll:
@pytest.mark.asyncio
async def test_scroll_start(self, client):
resp = await client.post(
"/api/scroll/start",
json={"mode": "target", "count": 5},
)
assert resp.status_code == 200
data = resp.json()
assert "success_radius" in data
assert "target_scrollY" in data
assert "direction" in data
assert data["success_radius"] == 80 # target mode
@pytest.mark.asyncio
async def test_scroll_skip(self, client):
# Start first
await client.post(
"/api/scroll/start",
json={"mode": "precise", "count": 3},
)
resp = await client.post(
"/api/scroll/skip",
json={"current_scrollY": 2000},
)
assert resp.status_code == 200
data = resp.json()
assert "target_scrollY" in data
assert "direction" in data
@pytest.mark.asyncio
async def test_scroll_status(self, client):
resp = await client.get("/api/scroll/status")
assert resp.status_code == 200
data = resp.json()
assert "trace_count" in data
assert "model_trained" in data

286
tests/tools/test_trainer.py Normal file
View File

@@ -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",
)