tests/test_model_support.py

from __future__ import annotations

from types import SimpleNamespace

import pytest

pytest.importorskip("torch")

from app.core.model_support import describe_model_support


def test_model_support_accepts_transformer_h_layout() -> None:
    layer = SimpleNamespace(attn=object())
    model = SimpleNamespace(
        transformer=SimpleNamespace(h=[layer, layer]),
        config=SimpleNamespace(_attn_implementation="eager"),
    )

    support = describe_model_support(model)

    assert support.supports_attribution is True
    assert support.layer_path == "transformer.h"
    assert support.attention_attr == "attn"


def test_model_support_rejects_non_eager_attention() -> None:
    layer = SimpleNamespace(self_attn=object())
    model = SimpleNamespace(
        model=SimpleNamespace(layers=[layer]),
        config=SimpleNamespace(_attn_implementation="sdpa"),
    )

    support = describe_model_support(model)

    assert support.supports_attribution is False
    assert "eager" in (support.reason or "")


def test_model_support_accepts_hrm_layout() -> None:
    layer = SimpleNamespace(attention=object())
    model = SimpleNamespace(
        model=SimpleNamespace(h_stack=[layer, layer]),
        config=SimpleNamespace(_attn_implementation="eager"),
    )

    support = describe_model_support(model)

    assert support.supports_attribution is True
    assert support.layer_path == "model.h_stack"
    assert support.attention_attr == "attention"