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, create_engine

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_hydrating_a_property_with_overrides_stays_on_one_connection(
    db_engine: Engine,
) -> None:
    """Resolving Landlord Overrides during re-hydration must read on the unit's
    OWN session, not open a second one. The modelling_e2e Lambda runs on a
    single-connection pool (``pool_size=1, max_overflow=0``); a second session
    opened while the unit's read transaction is still open checks out a second
    connection and deadlocks ("QueuePool limit of size 1 overflow 0 reached,
    connection timed out"). This pins that the unit holds exactly one connection
    while it hydrates a Property that has overrides.
    """
    # Arrange — seed a property + EPC + an override (committed reference data).
    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()

    # A pool that admits exactly one connection and fails fast (not after 30s) if
    # a second is requested — the production modelling_e2e shape.
    single_connection = create_engine(
        db_engine.url, pool_size=1, max_overflow=0, pool_timeout=2
    )
    try:
        # Act — hydrate through the unit; this resolves the override.
        with PostgresUnitOfWork(_session_factory(single_connection)) as uow:
            prop = uow.property.get(property_id)
    finally:
        single_connection.dispose()

    # Assert — reached here without a QueuePool timeout, and the override folded:
    # cavity (4) → solid brick (3) / internal (3).
    main = next(
        part
        for part in prop.effective_epc.sap_building_parts
        if part.identifier is BuildingPartIdentifier.MAIN
    )
    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