"""モデル生成のテスト（SPEC_02「モデル」: timm・出力 5 クラス）.

事前学習の重みはダウンロードしない（pretrained=False のみ）．
"""

from __future__ import annotations

import torch

from tias_judge_model.models.factory import count_parameters, create_model


def test_create_model_outputs_five_classes() -> None:
    torch.manual_seed(0)
    model = create_model("resnet18", pretrained=False).eval()
    with torch.no_grad():
        out = model(torch.zeros(2, 3, 32, 32))
    assert out.shape == (2, 5)


def test_count_parameters_matches_numel() -> None:
    model = create_model("resnet18", pretrained=False)
    assert count_parameters(model) == sum(p.numel() for p in model.parameters())
    assert count_parameters(model) > 0