87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
"""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()
|