ModelTrainEval/tests/test_config.py

from pathlib import Path

import pytest

pytest.importorskip("yaml")

from utils.config import Config, apply_overrides, deep_merge, load_config, parse_override  # noqa: E402


def test_deep_merge_is_recursive():
    merged = deep_merge({"a": {"b": 1, "c": 2}}, {"a": {"c": 3}, "d": 4})
    assert merged == {"a": {"b": 1, "c": 3}, "d": 4}


def test_parse_override_literals():
    assert parse_override("train.lr=1e-3") == ("train.lr", 1e-3)
    assert parse_override("train.amp=False") == ("train.amp", False)
    assert parse_override("model.backbone.name=tiny_conv") == ("model.backbone.name", "tiny_conv")
    assert parse_override("simclr.feature_stages=[5, 6]") == ("simclr.feature_stages", [5, 6])


def test_parse_override_accepts_yaml_style_scalars():
    """配置文件里写 null/true/false，命令行上也得能这么写。

    回归测试：只用 ast.literal_eval 时 `false` 会变成字符串 "false"，
    而非空字符串是真值——`--set xxx.enabled=false` 反而把开关打开了。
    """
    assert parse_override("a=null") == ("a", None)
    assert parse_override("a=~") == ("a", None)
    assert parse_override("a=true") == ("a", True)
    assert parse_override("a=false") == ("a", False)
    assert parse_override("a=false")[1] is False  # 不是字符串

    # 路径、模型名这类仍应保持字符串
    assert parse_override("a=resnet18") == ("a", "resnet18")
    assert parse_override("a=runs/x/best.pth") == ("a", "runs/x/best.pth")


def test_apply_overrides_creates_missing_nodes():
    cfg = apply_overrides({}, ["a.b.c=1"])
    assert cfg == {"a": {"b": {"c": 1}}}


def test_load_config_with_base(tmp_path):
    (tmp_path / "base.yaml").write_text("train:\n  lr: 0.1\n  epochs: 10\n", encoding="utf-8")
    (tmp_path / "child.yaml").write_text(
        "_base_: ./base.yaml\ntrain:\n  lr: 0.01\nwork_dir: runs/x\n", encoding="utf-8"
    )

    cfg = load_config(tmp_path / "child.yaml", ["train.epochs=2"])
    assert isinstance(cfg, Config)
    assert cfg.train.lr == 0.01
    assert cfg.train.epochs == 2
    assert cfg.get_path("work_dir") == "runs/x"


def test_load_config_detects_cycle(tmp_path):
    (tmp_path / "a.yaml").write_text("_base_: ./b.yaml\n", encoding="utf-8")
    (tmp_path / "b.yaml").write_text("_base_: ./a.yaml\n", encoding="utf-8")
    with pytest.raises(ValueError, match="循环引用"):
        load_config(tmp_path / "a.yaml")


def test_project_configs_are_loadable():
    cfg = load_config("configs/train/baseline.yaml")
    assert cfg["model"]["name"] == "defect_net"
    assert cfg["data"]["train"]["dataset"]["name"] == "film_defect"
    assert len(cfg["classes"]) > 0


def test_replace_key_swaps_a_whole_subdict(tmp_path):
    """换组件时要能整块替换，否则旧组件的专有字段会残留下来传给新组件。"""
    (tmp_path / "base.yaml").write_text(
        "backbone:\n  name: tiny_conv\n  widths: [32, 64]\n  depths: [1, 2]\n  block: dw\n",
        encoding="utf-8",
    )
    (tmp_path / "child.yaml").write_text(
        "_base_: ./base.yaml\nbackbone:\n  _replace_: true\n  name: torchvision\n  arch: resnet18\n",
        encoding="utf-8",
    )

    cfg = load_config(tmp_path / "child.yaml")

    assert cfg["backbone"] == {"name": "torchvision", "arch": "resnet18"}
    assert "_replace_" not in cfg["backbone"]  # 标记本身不能漏进构造参数


def test_without_replace_key_subdicts_still_merge(tmp_path):
    (tmp_path / "base.yaml").write_text("backbone:\n  name: tiny_conv\n  block: dw\n", encoding="utf-8")
    (tmp_path / "child.yaml").write_text(
        "_base_: ./base.yaml\nbackbone:\n  block: res\n", encoding="utf-8"
    )

    cfg = load_config(tmp_path / "child.yaml")
    assert cfg["backbone"] == {"name": "tiny_conv", "block": "res"}


def test_registry_reports_incompatible_fields_actionably():
    """--set 只改 name 时旧字段会残留，报错要说清楚是哪些字段、该怎么办。"""
    import pytest

    from models.registry import BACKBONES

    with pytest.raises(TypeError, match=r"不接受的字段.*widths"):
        BACKBONES.build({"name": "torchvision", "arch": "resnet18", "widths": [32]})


@pytest.mark.parametrize(
    "config_path",
    sorted(str(p) for p in Path("configs").rglob("*.yaml")),
)
def test_every_shipped_config_builds_its_components(config_path):
    """仓库里每个配置都必须能真的构建出组件。

    回归测试：configs/train/resnet18.yaml 曾经无法使用——它同时继承 tiny_conv 与
    resnet18 的模型配置，合并后 backbone 混了两边的字段，构建时直接报错。
    """
    torch = pytest.importorskip("torch")
    from datasets.yolo import load_classes
    from models import build_model
    from models.registry import BACKBONES

    cfg = load_config(config_path)

    if "model" in cfg:
        if "classes" not in cfg:
            # configs/model/*.yaml 是片段，只能被 _base_ 组合，单独用没有 classes
            pytest.skip(f"{config_path} 是配置片段，不含 classes")
        model_cfg = dict(cfg["model"])
        model_cfg.setdefault("num_classes", len(load_classes(cfg["classes"])))
        # 单测里不下载 ImageNet 权重
        if model_cfg.get("backbone", {}).get("name") == "torchvision":
            model_cfg["backbone"] = {**model_cfg["backbone"], "pretrained": False}
        model = build_model(model_cfg)
        with torch.no_grad():
            out = model(torch.randn(1, 3, 64, 64))
        assert "heatmap" in out

    simclr_backbone = cfg.get_path("simclr.backbone")
    if simclr_backbone:
        backbone_cfg = dict(simclr_backbone)
        if backbone_cfg.get("name") == "torchvision":
            backbone_cfg["pretrained"] = False
        backbone = BACKBONES.build(backbone_cfg)
        assert backbone.out_channels > 0


def test_eval_config_matches_its_training_backbone():
    """评估配置构建出的骨架必须与训练配置一致，否则权重根本加载不上。"""
    for train_cfg, eval_cfg in [
        ("configs/train/baseline.yaml", "configs/eval/baseline.yaml"),
        ("configs/train/resnet18.yaml", "configs/eval/resnet18.yaml"),
    ]:
        assert load_config(train_cfg)["model"] == load_config(eval_cfg)["model"], eval_cfg