refactor(tests): split into tests/unit and tests/tools
This commit is contained in:
56
tests/tools/conftest.py
Normal file
56
tests/tools/conftest.py
Normal 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
|
||||
Reference in New Issue
Block a user