caption-test/tests/test_trainer.py

"""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}"
    )