"""Tests for src.training.trainer_new.Trainer. Scope: __init__ behaviour, backbone validation, ModelsConfig type union. Out of scope: actual training (requires GPU + datasets + model checkpoints). The Trainer class is designed to defer all heavy lifting (CUDA, model construction, dataset loading) to .train(); __init__ just stores the 6 cfg objects and zeros out runtime state. This makes it cheap to test. """ from __future__ import annotations from typing import get_args import pytest from src.conf.config_loader import load_all_configs from src.conf.models_dinov3_conf import DINOv3ModelsConfig from src.conf.models_stripnet_conf import StripNetModelsConfig from src.training.trainer_new import ( ModelsConfig, Trainer, _SUPPORTED_BACKBONES, ) from conftest import DINOV3_PRESETS, STRIPNET_PRESETS # -- module-level constants ------------------------------------------------ def test_supported_backbones_is_frozenset() -> None: """_SUPPORTED_BACKBONES must be a frozenset (immutable, hashable).""" assert isinstance(_SUPPORTED_BACKBONES, frozenset) def test_supported_backbones_contents() -> None: """Exactly dinov3 and stripnet are supported in the current refactor. Sofia (v1/v71) is intentionally absent — see Trainer._validate_backbone for the rationale and the steps to add it later. """ assert _SUPPORTED_BACKBONES == frozenset({"dinov3", "stripnet"}) def test_models_config_union_contents() -> None: """ModelsConfig union mirrors _SUPPORTED_BACKBONES (dinov3 | stripnet).""" union_members = set(get_args(ModelsConfig)) assert union_members == {DINOv3ModelsConfig, StripNetModelsConfig} # -- _validate_backbone -------------------------------------------------- @pytest.mark.parametrize("backbone", ["dinov3", "stripnet"]) def test_validate_backbone_accepts_supported( path2cfg: str, backbone: str, ) -> None: """Supported backbones pass _validate_backbone silently. We use a real preset to build a valid Trainer — this also exercises the load_all_configs → Trainer(...) integration. """ preset = "gtauav_balanced" if backbone == "dinov3" else "gtauav_balanced_stripnet" cfgs = load_all_configs(path2cfg, preset) trainer = Trainer( pipeline_cfg=cfgs["pipeline"], hardware_cfg=cfgs["hardware"], training_cfg=cfgs["training"], tracking_cfg=cfgs["tracking"], models_common_cfg=cfgs["models_common"], models_cfg=cfgs["models"], ) # Must not raise. trainer._validate_backbone() @pytest.mark.parametrize("bad_backbone", ["sofia_v1", "sofia_v71", "mistral_42b", ""]) def test_validate_backbone_rejects_unsupported( path2cfg: str, bad_backbone: str, ) -> None: """Unsupported backbones (incl. sofia) raise NotImplementedError, not ImportError. The user must get a clear, actionable message — not a stack trace from a missing module. """ cfgs = load_all_configs(path2cfg, "gtauav_balanced") trainer = Trainer( pipeline_cfg=cfgs["pipeline"], hardware_cfg=cfgs["hardware"], training_cfg=cfgs["training"], tracking_cfg=cfgs["tracking"], models_common_cfg=cfgs["models_common"], models_cfg=cfgs["models"], ) # Tamper with backbone — simulate what would happen if config_loader # were extended to accept sofia presets. trainer.models_common_cfg.backbone = bad_backbone with pytest.raises(NotImplementedError) as excinfo: trainer._validate_backbone() # Error message must mention the offending backbone name and what's supported. msg = str(excinfo.value) assert bad_backbone in msg or repr(bad_backbone) in msg assert "dinov3" in msg assert "stripnet" in msg # -- Trainer.__init__ smoke tests ------------------------------------------ @pytest.mark.parametrize("preset_name", DINOV3_PRESETS + STRIPNET_PRESETS) def test_trainer_init_with_real_preset(path2cfg: str, preset_name: str) -> None: """Trainer(...) instantiates from every real preset's loaded cfgs. Heavy work (CUDA, model build, dataset open) is deferred to .train(); __init__ only stores cfgs and zeros runtime state, so this is cheap and GPU-free. """ cfgs = load_all_configs(path2cfg, preset_name) trainer = Trainer( pipeline_cfg=cfgs["pipeline"], hardware_cfg=cfgs["hardware"], training_cfg=cfgs["training"], tracking_cfg=cfgs["tracking"], models_common_cfg=cfgs["models_common"], models_cfg=cfgs["models"], ) # Cfgs are stored as-is. assert trainer.pipeline_cfg is cfgs["pipeline"] assert trainer.hardware_cfg is cfgs["hardware"] assert trainer.training_cfg is cfgs["training"] assert trainer.tracking_cfg is cfgs["tracking"] assert trainer.models_common_cfg is cfgs["models_common"] assert trainer.models_cfg is cfgs["models"] def test_trainer_init_zeros_runtime_state(path2cfg: str) -> None: """All runtime fields are None / 0 / [] before .train() is called.""" cfgs = load_all_configs(path2cfg, "gtauav_balanced") trainer = Trainer( pipeline_cfg=cfgs["pipeline"], hardware_cfg=cfgs["hardware"], training_cfg=cfgs["training"], tracking_cfg=cfgs["tracking"], models_common_cfg=cfgs["models_common"], models_cfg=cfgs["models"], ) # None-typed runtime fields. for attr in ( "output_dir", "full_config", "tracker", "csv_logger", "model", "loss_fn", "neg_bank", "optimizer", "scheduler", "scaler", "train_ds", "test_ds", "train_eval_ds", "train_loader", "test_loader", "train_eval_loader", "batch_sampler", "emb_cache", "profiler", "resume_ckpt", ): assert getattr(trainer, attr) is None, ( f"trainer.{attr} should be None before .train(), " f"got {type(getattr(trainer, attr)).__name__}" ) # Counter / loop state initialized to identity values. assert trainer.start_epoch == 0 assert trainer.global_step == 0 assert trainer.best_r1 == 0.0 assert trainer.history == [] assert trainer.steps_per_epoch == 0 # -- Trainer.train end-to-end signature ------------------------------------ def test_trainer_train_method_exists_and_takes_no_args() -> None: """Trainer.train() takes only `self` — main.py calls trainer.train().""" import inspect sig = inspect.signature(Trainer.train) params = [p for p in sig.parameters.values() if p.name != "self"] assert params == [], ( f"Trainer.train() must take only self; got extra params: {params}" )