Model/tests/repositories/test_unit_of_work.py

from __future__ import annotations

import json
from collections.abc import Callable
from pathlib import Path
from typing import Any

import pytest
from sqlalchemy import Engine
from sqlmodel import Session

from datatypes.epc.domain.epc import Epc
from datatypes.epc.domain.epc_property_data import (
    BuildingPartIdentifier,
    EpcPropertyData,
)
from datatypes.epc.domain.mapper import EpcPropertyDataMapper
from domain.property_baseline.property_baseline_performance import PropertyBaselinePerformance
from domain.property_baseline.performance import Performance
from infrastructure.postgres.property_override_table import PropertyOverrideRow
from infrastructure.postgres.property_table import PropertyRow
from repositories.epc.epc_postgres_repository import EpcPostgresRepository
from repositories.plan.plan_repository import PlanRepository
from repositories.postgres_unit_of_work import PostgresUnitOfWork
from repositories.product.product_repository import ProductRepository
from repositories.scenario.scenario_repository import ScenarioRepository
from repositories.spatial.spatial_repository import SpatialRepository

_JSON_SAMPLES = Path(__file__).resolve().parents[2] / "backend/epc_api/json_samples"


def _session_factory(db_engine: Engine) -> Callable[[], Session]:
    return lambda: Session(db_engine)


def _epc() -> EpcPropertyData:
    raw: dict[str, Any] = json.loads(
        (_JSON_SAMPLES / "RdSAP-Schema-21.0.0" / "epc.json").read_text()
    )
    return EpcPropertyDataMapper.from_api_response(raw)


def _baseline() -> PropertyBaselinePerformance:
    perf = Performance(
        sap_score=72, epc_band=Epc.C, co2_emissions=1.8, primary_energy_intensity=180
    )
    return PropertyBaselinePerformance(
        lodged=perf,
        effective=perf,
        rebaseline_reason="none",
        space_heating_kwh=5000.0,
        water_heating_kwh=2000.0,
    )


def test_committed_work_is_visible_to_a_later_unit(db_engine: Engine) -> None:
    # Arrange
    new_unit = lambda: PostgresUnitOfWork(_session_factory(db_engine))
    baseline = _baseline()

    # Act
    with new_unit() as uow:
        uow.property_baseline.save(baseline, property_id=10)
        uow.commit()

    # Assert — a fresh unit reads back what the first one committed.
    with new_unit() as uow:
        loaded = uow.property_baseline.get_for_property(10)
    assert loaded == baseline


def test_an_exception_in_the_block_rolls_the_batch_back(db_engine: Engine) -> None:
    # Arrange
    new_unit = lambda: PostgresUnitOfWork(_session_factory(db_engine))

    # Act — a property mid-batch raises after a write but before commit.
    with pytest.raises(RuntimeError, match="boom"):
        with new_unit() as uow:
            uow.property_baseline.save(_baseline(), property_id=10)
            raise RuntimeError("boom")

    # Assert — nothing from the aborted batch is persisted.
    with new_unit() as uow:
        assert uow.property_baseline.get_for_property(10) is None


def test_unit_exposes_the_modelling_repos_bound_to_its_session(
    db_engine: Engine,
) -> None:
    # Arrange / Act
    with PostgresUnitOfWork(_session_factory(db_engine)) as uow:
        # Assert — the Modelling stage reads Scenario + Product and writes Plan
        # through the same unit (ADR-0017).
        assert isinstance(uow.scenario, ScenarioRepository)
        assert isinstance(uow.product, ProductRepository)
        assert isinstance(uow.plan, PlanRepository)


def test_unit_exposes_the_spatial_cache_repo_bound_to_its_session(
    db_engine: Engine,
) -> None:
    # Arrange / Act
    with PostgresUnitOfWork(_session_factory(db_engine)) as uow:
        # Assert — Ingestion writes the OS spatial reference cache through the
        # same unit it persists the EPC/solar with (ADR-0020).
        assert isinstance(uow.spatial, SpatialRepository)


def test_unit_hydrates_a_property_with_its_landlord_overrides_folded(
    db_engine: Engine,
) -> None:
    """A Property re-hydrated through the unit folds its Landlord Overrides onto
    the Effective EPC (ADR-0032) — the same overlay the live modelling path
    applies. Without this the Baseline Performance the orchestrator scores off
    ``uow.property`` diverges from the Plan, which does apply the overrides.
    """
    # Arrange — a lodged EPC (cavity main wall, 4) plus a solid-brick / internal
    # wall override, persisted through the unit.
    with PostgresUnitOfWork(_session_factory(db_engine)) as uow:
        row = PropertyRow(portfolio_id=1, postcode="A0 0AA", address="1 St", uprn=1)
        uow._session.add(row)  # pyright: ignore[reportPrivateUsage]
        uow._session.flush()  # pyright: ignore[reportPrivateUsage]
        property_id = row.id
        assert property_id is not None
        EpcPostgresRepository(uow._session).save(  # pyright: ignore[reportPrivateUsage]
            _epc(), property_id=property_id
        )
        uow._session.add(  # pyright: ignore[reportPrivateUsage]
            PropertyOverrideRow(
                property_id=property_id,
                portfolio_id=1,
                building_part=0,
                override_component="wall_type",
                override_value="Solid brick, with internal insulation",
                original_spreadsheet_description="solid brick, insulated",
            )
        )
        uow.commit()

    # Act — re-hydrate through a fresh unit.
    with PostgresUnitOfWork(_session_factory(db_engine)) as uow:
        prop = uow.property.get(property_id)

    main = next(
        part
        for part in prop.effective_epc.sap_building_parts
        if part.identifier is BuildingPartIdentifier.MAIN
    )

    # Assert — the override is folded: cavity (4) → solid brick (3) / internal (3).
    assert main.wall_construction == 3
    assert main.wall_insulation_type == 3


def test_leaving_the_block_without_commit_persists_nothing(db_engine: Engine) -> None:
    # Arrange
    new_unit = lambda: PostgresUnitOfWork(_session_factory(db_engine))

    # Act — write but never commit.
    with new_unit() as uow:
        uow.property_baseline.save(_baseline(), property_id=10)

    # Assert
    with new_unit() as uow:
        assert uow.property_baseline.get_for_property(10) is None