tests/test_llm_fake.py

from __future__ import annotations

import pytest

from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import (
    NonRetryableModelError,
    RetryableModelError,
    StructuredModelResponseError,
)
from gcmd_classifier.llm import FakeModelClient, ModelRequest, ModelStage
from gcmd_classifier.llm.base import RetryPolicy, generate_with_retries
from gcmd_classifier.llm.schemas import TermResponse, TopicResponse, VariableResponse


def _request(stage: ModelStage, schema: type) -> ModelRequest:
    settings = ModelSettings()
    return ModelRequest.from_settings(
        stage=stage,
        prompt="prompt",
        response_schema=schema,
        settings=settings,
    )


def _decision(candidate_id: str = "candidate-1") -> dict:
    return {
        "candidate_id": candidate_id,
        "confidence": 0.8,
        "evidence": "The article supports this candidate.",
        "support_type": "explicit",
    }


def test_fake_model_returns_typed_topic_response() -> None:
    client = FakeModelClient([{"selected": [_decision()], "ambiguous_alternatives": []}])

    response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse))

    assert isinstance(response.parsed, TopicResponse)
    assert response.parsed.selected[0].candidate_id == "candidate-1"
    assert response.provider == "fake"


def test_fake_model_returns_typed_term_response() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("term-1")],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    response = client.generate_structured(_request(ModelStage.TERM, TermResponse))

    assert isinstance(response.parsed, TermResponse)
    assert response.parsed.selected[0].candidate_id == "term-1"


def test_fake_model_returns_typed_variable_response() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision("variable-1")],
                "stop_at_parent": False,
                "ambiguous_alternatives": [],
            }
        ]
    )

    response = client.generate_structured(_request(ModelStage.VARIABLE, VariableResponse))

    assert isinstance(response.parsed, VariableResponse)
    assert response.parsed.selected[0].candidate_id == "variable-1"


def test_fake_model_can_return_no_selection() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [],
                "ambiguous_alternatives": [],
                "no_selection_reason": "No supplied Topic is supported.",
            }
        ]
    )

    response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse))

    assert response.parsed.selected == []
    assert response.parsed.no_selection_reason == "No supplied Topic is supported."


def test_fake_model_records_requests() -> None:
    client = FakeModelClient([{"selected": [], "ambiguous_alternatives": []}])
    request = _request(ModelStage.TOPIC, TopicResponse)

    client.generate_structured(request)

    assert client.requests == [request]


def test_fake_model_can_return_invalid_candidate_ids_for_later_validation() -> None:
    client = FakeModelClient([{"selected": [_decision("not-a-supplied-id")]}])

    response = client.generate_structured(_request(ModelStage.TOPIC, TopicResponse))

    assert response.parsed.selected[0].candidate_id == "not-a-supplied-id"


def test_fake_model_rejects_contradictory_structured_response() -> None:
    client = FakeModelClient(
        [
            {
                "selected": [_decision()],
                "stop_at_parent": True,
                "stop_reason": "Contradictory response.",
                "ambiguous_alternatives": [],
            }
        ]
    )

    with pytest.raises(StructuredModelResponseError):
        client.generate_structured(_request(ModelStage.TERM, TermResponse))


def test_fake_model_can_raise_retryable_errors() -> None:
    client = FakeModelClient([RetryableModelError("temporary failure")])

    with pytest.raises(RetryableModelError):
        client.generate_structured(_request(ModelStage.TOPIC, TopicResponse))


def test_fake_model_can_raise_non_retryable_errors() -> None:
    client = FakeModelClient([NonRetryableModelError("permanent failure")])

    with pytest.raises(NonRetryableModelError):
        client.generate_structured(_request(ModelStage.TOPIC, TopicResponse))


def test_fake_model_can_fail_once_then_succeed_with_retry_wrapper() -> None:
    client = FakeModelClient(
        [
            RetryableModelError("temporary failure"),
            {"selected": [_decision()], "ambiguous_alternatives": []},
        ]
    )

    response = generate_with_retries(
        client,
        _request(ModelStage.TOPIC, TopicResponse),
        RetryPolicy(max_retries=1),
    )

    assert response.retry_count == 1
    assert len(client.requests) == 2
    assert response.parsed.selected[0].candidate_id == "candidate-1"


def test_fake_model_can_exhaust_retries() -> None:
    client = FakeModelClient(
        [RetryableModelError("one"), RetryableModelError("two"), RetryableModelError("three")]
    )

    with pytest.raises(Exception) as exc_info:
        generate_with_retries(
            client,
            _request(ModelStage.TOPIC, TopicResponse),
            RetryPolicy(max_retries=2),
        )

    assert exc_info.value.__class__.__name__ == "ModelRetriesExhaustedError"
    assert len(client.requests) == 3