Model/tests/domain/data_transformation/test_guarded_column_classifier.py

from __future__ import annotations

from enum import Enum
from typing import Optional

from domain.data_transformation.column_classifier import ColumnClassifier
from domain.data_transformation.guarded_column_classifier import (
    GuardedColumnClassifier,
)


class _Category(Enum):
    GUARDED = "guarded"
    FALLBACK = "fallback"
    UNKNOWN = "unknown"


class _RecordingFallback(ColumnClassifier[_Category]):
    """Stands in for the LLM: records what it was asked and maps everything it
    sees to FALLBACK, so the test can see which descriptions reached it."""

    def __init__(self) -> None:
        self.asked: set[str] = set()

    def classify(self, descriptions: set[str]) -> dict[str, _Category]:
        self.asked = set(descriptions)
        return {d: _Category.FALLBACK for d in descriptions}


def _guard(description: str) -> Optional[_Category]:
    return _Category.GUARDED if description == "marker" else None


def test_guard_hits_win_and_misses_fall_through_to_the_fallback() -> None:
    # Arrange
    fallback = _RecordingFallback()
    classifier = GuardedColumnClassifier(guard=_guard, fallback=fallback)

    # Act
    result = classifier.classify({"marker", "other"})

    # Assert — the guarded description takes the guard's member and never reaches
    # the fallback; the unrecognised one is resolved by the fallback.
    assert result == {"marker": _Category.GUARDED, "other": _Category.FALLBACK}
    assert fallback.asked == {"other"}