Model/tests/domain/epc/test_main_heating_guard.py

from __future__ import annotations

import pytest

from domain.epc.property_overrides.main_heating_guard import main_heating_guard
from domain.epc.property_overrides.main_heating_system_type import (
    MainHeatingSystemType,
)


def test_guard_resolves_high_heat_retention_storage() -> None:
    # Arrange — the structured phrasing the LLM funnelled into "old storage" (401).
    description = "Electric Storage Systems: High heat retention storage heaters"

    # Act
    result = main_heating_guard(description)

    # Assert — HHRSH has its own archetype (SAP 409), never old storage (ADR-0044).
    assert result is MainHeatingSystemType.ELECTRIC_STORAGE_HIGH_HEAT_RETENTION


@pytest.mark.parametrize(
    ("description", "expected"),
    [
        # Electric underfloor dumped into Direct-acting / old storage (ADR-0046).
        (
            "Electric Underfloor Heating: In concrete slab (off-peak only)",
            MainHeatingSystemType.ELECTRIC_UNDERFLOOR_SLAB_OFF_PEAK,
        ),
        (
            "Electric Underfloor Heating: Integrated (storage+direct-acting) (off peak)",
            MainHeatingSystemType.ELECTRIC_UNDERFLOOR_INTEGRATED,
        ),
        (
            "Electric Underfloor Heating: In screed above insulation (standard or off peak)",
            MainHeatingSystemType.ELECTRIC_UNDERFLOOR_SCREED,
        ),
    ],
)
def test_guard_resolves_electric_underfloor(
    description: str, expected: MainHeatingSystemType
) -> None:
    # Act
    result = main_heating_guard(description)

    # Assert
    assert result is expected


@pytest.mark.parametrize(
    ("description", "expected"),
    [
        # The Gas CPSU dumping ground (ADR-0045): solid-fuel room heaters and
        # electric ("NA" type) boilers the LLM sent to Gas CPSU (120, mains gas).
        (
            "Solid fuel room heaters: Open fire in grate",
            MainHeatingSystemType.SOLID_FUEL_ROOM_HEATER_OPEN_FIRE,
        ),
        (
            "Solid fuel room heaters: Open fire with back boiler (no radiators)",
            MainHeatingSystemType.SOLID_FUEL_ROOM_HEATER_OPEN_FIRE_BACK_BOILER,
        ),
        (
            "Solid fuel room heaters: Closed room heater with boiler (no radiators)",
            MainHeatingSystemType.SOLID_FUEL_ROOM_HEATER_CLOSED_WITH_BOILER,
        ),
        # "NA" is the boiler-type value meaning "not a gas combi/regular/CPSU" —
        # i.e. an electric boiler.
        ("Boiler: A rated NA", MainHeatingSystemType.ELECTRIC_BOILER),
    ],
)
def test_guard_resolves_the_gas_cpsu_dumping_ground(
    description: str, expected: MainHeatingSystemType
) -> None:
    # Act
    result = main_heating_guard(description)

    # Assert
    assert result is expected


@pytest.mark.parametrize(
    "description",
    [
        # Other storage subtypes the LLM already classifies correctly — the guard
        # only rescues the HHRSH gap, leaving the rest to the fallback classifier.
        "Electric Storage Systems: Old (large volume) storage heaters",
        "Electric Storage Systems: Modern (slimline) storage heaters",
        "Electric Storage Systems: Fan storage heaters",
        # Genuinely-gas boilers the LLM already classifies correctly (not "NA").
        "Boiler: A rated Combi",
        "Boiler: C rated Regular Boiler",
        # A plain closed room heater (no boiler) is the existing 633 archetype —
        # the LLM keeps it; only "…with boiler" (634) is guarded.
        "Solid fuel room heaters: Closed room heater (no boiler)",
        # Unrelated / varied phrasings are the LLM's job.
        "Gas boiler",
        "",
    ],
)
def test_guard_defers_non_hhrsh_descriptions_to_the_llm(description: str) -> None:
    # Act
    result = main_heating_guard(description)

    # Assert
    assert result is None