junte/Model
1
0
Fork 0
mirror of https://github.com/Hestia-Homes/Model.git synced 2026-07-12 13:29:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
Model/tests/applications/modelling_e2e/test_handler.py
Jun-te Kim b1ff711260 perf(modelling_e2e): batch SubTask bookkeeping to stop per-property writes 
Even after batching the data writes, the handler still wrote to the DB per
property through the orchestrator's SubTask bookkeeping: create + start +
complete each self-committed, and _cascade re-listed every sibling and re-saved
the parent on every transition — ~5 writes per property plus an O(N^2) cascade.

- TaskOrchestrator.run_subtasks: create all children in one INSERT, run each
  (failures isolated per child), then persist all terminal states in one bulk
  save and cascade the parent once. Children go WAITING -> terminal; the
  transient IN_PROGRESS row is never written.
- SubTaskRepository.create_many / save_many (bulk INSERT / bulk fetch + update).
- _cascade short-circuits when the Task is already FAILED (terminal) — skips the
  sibling roll-up entirely.
- modelling_e2e handler fans out via run_subtasks instead of per-property
  create_child_subtask + run_subtask.

Per N-property batch the SubTask bookkeeping drops from ~5N writes + an O(N^2)
cascade to ~2 writes + 1 cascade.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-24 19:26:42 +00:00

1124 lines
44 KiB
Python
Raw Blame History

"""Tests for the modelling_e2e Lambda handler.
Tests exercise the handler's external behaviour through handler.__wrapped__,
patching I/O boundaries (EPC client, DB reads, UoW) so no real DB or network
is needed. One test per distinct behaviour path.
"""
from __future__ import annotations
from contextlib import ExitStack
from typing import Any, Iterator, Optional
from unittest.mock import MagicMock, call, patch
from uuid import UUID, uuid4
import pytest
from pydantic import ValidationError
from applications.modelling_e2e.modelling_e2e_trigger_body import (
ModellingE2ETriggerBody,
)
from domain.tasks.subtasks import SubTask
PROPERTY_ID = 12345
UPRN = 987654321
POSTCODE = "SW1A 1AA"
PORTFOLIO_ID = 100
SCENARIO_ID = 200
_ENV = {
"OPEN_EPC_API_TOKEN": "test-token",
"DATA_BUCKET": "test-bucket",
"GOOGLE_SOLAR_API_KEY": "test-solar-key",
}
_BODY = {
"property_ids": [PROPERTY_ID],
"portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID,
"no_solar": True,
"dry_run": False,
}
def _mock_orchestrator() -> MagicMock:
"""Mock TaskOrchestrator whose run_subtasks creates one SubTask per input and
runs work for each, isolating per-item failures (mirroring the real method)."""
mock = MagicMock()
def _run_subtasks(
parent_task_id: UUID,
inputs_per_subtask: list[dict[str, Any]],
work: Any,
**kwargs: Any,
) -> list[Any]:
results: list[Any] = []
for inputs in inputs_per_subtask:
subtask = SubTask.create(task_id=parent_task_id, inputs=inputs)
try:
results.append(work(subtask))
except Exception: # noqa: BLE001 — siblings continue, as in real impl
results.append(None)
return results
mock.run_subtasks.side_effect = _run_subtasks
return mock
def _call_handler(
body: dict[str, Any], orchestrator: Optional[MagicMock] = None
) -> Any:
from applications.modelling_e2e.handler import handler
orch = orchestrator if orchestrator is not None else _mock_orchestrator()
return handler.__wrapped__(body, None, orch, uuid4()) # type: ignore[attr-defined]
def _engine_mock(
property_ids: list[int],
uprns: list[int],
postcodes: list[str],
) -> MagicMock:
"""Mock engine whose connect() returns one (id, uprn, postcode) row set —
the handler reads all three columns in a single query."""
mock_engine = MagicMock()
mock_conn = mock_engine.connect.return_value.__enter__.return_value
property_result = MagicMock()
property_result.fetchall.return_value = list(
zip(property_ids, uprns, postcodes)
)
mock_conn.execute.return_value = property_result
return mock_engine
def _plan_mock() -> MagicMock:
plan = MagicMock()
plan.measures = []
plan.cost_of_works = 0.0
return plan
# ---------------------------------------------------------------------------
# Fixture: clear module-level cohort cache between tests
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def _clear_cohort_cache() -> None:
import applications.modelling_e2e.handler as h
h._cohort_cache.clear()
h._nearby_cohort_cache.clear()
@pytest.fixture(autouse=True)
def _baseline_orchestrator() -> Iterator[MagicMock]:
"""Patch the Baseline orchestrator for every test — construction stays cheap
and ``.run`` is a no-op mock the baseline tests assert against."""
with patch(
"applications.modelling_e2e.handler.PropertyBaselineOrchestrator"
) as orchestrator:
yield orchestrator
# ---------------------------------------------------------------------------
# Trigger body validation
# ---------------------------------------------------------------------------
def test_trigger_body_requires_property_ids() -> None:
"""property_ids (plural) must be provided and must be a list of ints."""
# Arrange
body = {
"property_ids": [1, 2],
"portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID,
}
# Act
result = ModellingE2ETriggerBody.model_validate(body)
# Assert
assert result.property_ids == [1, 2]
def test_trigger_body_rejects_missing_property_ids() -> None:
with pytest.raises(ValidationError):
ModellingE2ETriggerBody.model_validate(
{"portfolio_id": PORTFOLIO_ID, "scenario_id": SCENARIO_ID}
)
# ---------------------------------------------------------------------------
# Child SubTask creation
# ---------------------------------------------------------------------------
def test_handler_creates_one_child_subtask_per_property_id() -> None:
"""Handler creates exactly N child SubTasks, one per property in the batch,
each recording its property_id in inputs."""
# Arrange
pid1, pid2, pid3 = 111, 222, 333
mock_engine = _engine_mock(
[pid1, pid2, pid3], [1001, 1002, 1003], [POSTCODE, POSTCODE, POSTCODE]
)
mock_orch = _mock_orchestrator()
task_id = uuid4()
with ExitStack() as stack:
stack.enter_context(patch("applications.modelling_e2e.handler.os.environ", _ENV))
stack.enter_context(
patch("applications.modelling_e2e.handler._get_engine", return_value=mock_engine)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(patch("applications.modelling_e2e.handler.GeospatialS3Repository"))
stack.enter_context(patch("applications.modelling_e2e.handler.GoogleSolarApiClient"))
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler._solar_insights_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"))
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch("applications.modelling_e2e.handler.run_modelling", return_value=_plan_mock())
)
MockUoW = stack.enter_context(patch("applications.modelling_e2e.handler.PostgresUnitOfWork"))
mock_uow = MagicMock()
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
from applications.modelling_e2e.handler import handler
handler.__wrapped__( # type: ignore[attr-defined]
{"property_ids": [pid1, pid2, pid3], "portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID, "no_solar": True, "dry_run": False},
None, mock_orch, task_id,
)
# Assert — the batch is fanned out in ONE run_subtasks call, with one input
# per property (each recording its property_id) under the same parent task.
mock_orch.run_subtasks.assert_called_once()
args = mock_orch.run_subtasks.call_args.args
assert args[0] == task_id
inputs_per_subtask = args[1]
assert [i["property_id"] for i in inputs_per_subtask] == [pid1, pid2, pid3]
# ---------------------------------------------------------------------------
# Lodged EPC path
# ---------------------------------------------------------------------------
def test_lodged_epc_path_saves_epc_plan_and_marks_modelled(
_baseline_orchestrator: MagicMock,
) -> None:
"""When get_by_uprn returns an EPC the handler saves it, saves the plan,
marks the property as modelled — all inside one UoW per property — and
re-establishes its Baseline Performance after the plan commits."""
# Arrange
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
mock_epc = MagicMock()
mock_plan = _plan_mock()
mock_uow = MagicMock()
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = mock_epc
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._spatial_for", return_value=None
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._solar_insights_for",
return_value=None,
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.overlays_from", return_value=[]
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.Session")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=mock_plan,
)
)
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
_call_handler(_BODY)
# Assert — EPC saved (lodged path), plan saved, property marked modelled
mock_uow.epc.save.assert_called_once_with(
mock_epc, property_id=PROPERTY_ID, portfolio_id=PORTFOLIO_ID
)
mock_uow.plan.save.assert_called_once()
mock_uow.property.mark_modelled.assert_called_once_with(
PROPERTY_ID, has_recommendations=False
)
mock_uow.commit.assert_called_once()
# Baseline Performance is re-established for the lodged property.
_baseline_orchestrator.return_value.run.assert_called_once_with([PROPERTY_ID])
def test_skipped_cohort_certs_do_not_prevent_plan_being_saved() -> None:
"""Cohort certs the mapper can't consume are skipped so prediction can
proceed — the property's plan is committed before the handler returns."""
from repositories.comparable_properties.epc_comparable_properties_repository import (
SkippedCohortCert,
)
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
mock_plan = _plan_mock()
mock_uow = MagicMock()
skipped = [
SkippedCohortCert(
certificate_number="8257-7539-1649-0633-4992",
error="ValueError: RdSapSchema17_1: missing required field 'window'",
)
]
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._solar_insights_for",
return_value=None,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch(
"applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"
)
)
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=mock_plan,
)
)
# The repo accumulated a skipped (unmappable) cohort cert during the run.
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcComparablePropertiesRepository")
).return_value.skipped = skipped
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act — no raise even with a skipped cert
_call_handler(_BODY)
# Assert — plan committed despite the skipped cert
mock_uow.plan.save.assert_called_once()
mock_uow.commit.assert_called_once()
def test_skipped_cohort_certs_are_logged_and_handler_does_not_raise() -> None:
"""Unmappable cohort certs are logged but do not raise — the coordinator
SubTask completes normally and the property plan is committed."""
from repositories.comparable_properties.epc_comparable_properties_repository import (
SkippedCohortCert,
)
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
mock_plan = _plan_mock()
mock_uow = MagicMock()
skipped = [
SkippedCohortCert(
certificate_number="8257-7539-1649-0633-4992",
error="ValueError: RdSapSchema17_1: missing required field 'window'",
)
]
with ExitStack() as stack:
stack.enter_context(patch("applications.modelling_e2e.handler.os.environ", _ENV))
stack.enter_context(
patch("applications.modelling_e2e.handler._get_engine", return_value=mock_engine)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(patch("applications.modelling_e2e.handler.GeospatialS3Repository"))
stack.enter_context(patch("applications.modelling_e2e.handler.GoogleSolarApiClient"))
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler._solar_insights_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"))
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch("applications.modelling_e2e.handler.run_modelling", return_value=mock_plan)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcComparablePropertiesRepository")
).return_value.skipped = skipped
MockUoW = stack.enter_context(patch("applications.modelling_e2e.handler.PostgresUnitOfWork"))
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
mock_logger = stack.enter_context(
patch("applications.modelling_e2e.handler.logger")
)
# Act — must not raise even though cohort certs were skipped
_call_handler(_BODY)
# Assert — plan committed; skipped cert number surfaced in a log call
mock_uow.plan.save.assert_called_once()
mock_uow.commit.assert_called_once()
logged_messages = " ".join(
str(c.args) + str(c.kwargs) for c in mock_logger.info.call_args_list
)
assert "8257-7539-1649-0633-4992" in logged_messages
# ---------------------------------------------------------------------------
# EPC Prediction path
# ---------------------------------------------------------------------------
def test_prediction_path_saves_predicted_epc_plan_and_baseline(
_baseline_orchestrator: MagicMock,
) -> None:
"""When get_by_uprn returns None the handler synthesises an EPC via
prediction, persists it in the predicted slot (source='predicted'), saves the
plan, and re-establishes the Baseline from the predicted picture."""
# Arrange
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
mock_plan = _plan_mock()
mock_uow = MagicMock()
mock_predicted_epc = MagicMock()
# _predict_epc checks for a MAIN building part
from datatypes.epc.domain.epc_property_data import BuildingPartIdentifier
mock_part = MagicMock()
mock_part.identifier = BuildingPartIdentifier.MAIN
mock_predicted_epc.sap_building_parts = [mock_part]
mock_comparables = MagicMock()
mock_comparables.members = [MagicMock()] # non-empty cohort
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = None # no lodged EPC
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._spatial_for", return_value=None
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._solar_insights_for",
return_value=None,
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.overlays_from", return_value=[]
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
# Prediction infrastructure
from domain.epc_prediction.prediction_target import PredictionTargetAttributes
stack.enter_context(
patch(
"applications.modelling_e2e.handler.OverrideBackedPredictionAttributesReader"
)
).return_value.attributes_for.return_value = PredictionTargetAttributes(
property_type="2"
)
stack.enter_context(
patch("applications.modelling_e2e.handler.select_comparables")
).return_value = mock_comparables
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcPrediction")
).return_value.predict.return_value = mock_predicted_epc
stack.enter_context(
patch(
"applications.modelling_e2e.handler.EpcComparablePropertiesRepository"
)
).return_value.candidates_for.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.Session")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=mock_plan,
)
)
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
_call_handler(_BODY)
# Assert — predicted EPC persisted in the predicted slot, plan saved, baseline run
mock_uow.epc.save.assert_called_once_with(
mock_predicted_epc,
property_id=PROPERTY_ID,
portfolio_id=PORTFOLIO_ID,
source="predicted",
)
mock_uow.plan.save.assert_called_once()
mock_uow.commit.assert_called_once()
_baseline_orchestrator.return_value.run.assert_called_once_with([PROPERTY_ID])
# ---------------------------------------------------------------------------
# Prediction gate-out (empty cohort)
# ---------------------------------------------------------------------------
def test_empty_cohort_gates_property_out_without_saving() -> None:
"""When candidates_for returns an empty list the property cannot be predicted;
the failure is recorded in the child SubTask and the handler returns without
raising — no plan is saved."""
# Arrange
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
empty_comparables = MagicMock()
empty_comparables.members = [] # empty cohort → gate-out
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = None
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._spatial_for", return_value=None
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.overlays_from", return_value=[]
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
from domain.epc_prediction.prediction_target import PredictionTargetAttributes
stack.enter_context(
patch(
"applications.modelling_e2e.handler.OverrideBackedPredictionAttributesReader"
)
).return_value.attributes_for.return_value = PredictionTargetAttributes(
property_type="2"
)
stack.enter_context(
patch("applications.modelling_e2e.handler.select_comparables")
).return_value = empty_comparables
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcPrediction")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.EpcComparablePropertiesRepository"
)
).return_value.candidates_for.return_value = []
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides")
)
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
# Act — no raise; child SubTask is failed, handler continues
_call_handler(_BODY)
# UoW never entered — the property errored before the write block
MockUoW.return_value.__enter__.assert_not_called()
# ---------------------------------------------------------------------------
# Broadened cohort — sparse own postcode falls back to nearby postcodes
# ---------------------------------------------------------------------------
def test_empty_own_postcode_broadens_to_nearby_and_predicts() -> None:
"""When the property's own postcode holds no same-type comparables, the
handler broadens to the nearby-postcode cohort (candidates_near) and, finding
comparables there, synthesises the EPC and saves the plan."""
# Arrange
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
mock_plan = _plan_mock()
mock_uow = MagicMock()
mock_predicted_epc = MagicMock()
from datatypes.epc.domain.epc_property_data import BuildingPartIdentifier
mock_part = MagicMock()
mock_part.identifier = BuildingPartIdentifier.MAIN
mock_predicted_epc.sap_building_parts = [mock_part]
# First select_comparables (own postcode) is empty → broaden; the second
# (nearby cohort) finds comparables.
empty_comparables = MagicMock()
empty_comparables.members = []
found_comparables = MagicMock()
found_comparables.members = [MagicMock()]
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = None # no lodged EPC
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._solar_insights_for",
return_value=None,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
from domain.epc_prediction.prediction_target import PredictionTargetAttributes
stack.enter_context(
patch(
"applications.modelling_e2e.handler.OverrideBackedPredictionAttributesReader"
)
).return_value.attributes_for.return_value = PredictionTargetAttributes(
property_type="2"
)
MockRepo = stack.enter_context(
patch(
"applications.modelling_e2e.handler.EpcComparablePropertiesRepository"
)
)
MockRepo.return_value.candidates_for.return_value = []
MockRepo.return_value.candidates_near.return_value = [MagicMock()]
stack.enter_context(
patch(
"applications.modelling_e2e.handler.select_comparables",
side_effect=[empty_comparables, found_comparables],
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcPrediction")
).return_value.predict.return_value = mock_predicted_epc
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch(
"applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"
)
)
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=mock_plan,
)
)
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
_call_handler(_BODY)
# Assert — broadening fired, and the broadened cohort produced a saved plan
# with its predicted EPC persisted in the predicted slot.
MockRepo.return_value.candidates_near.assert_called_once()
mock_uow.epc.save.assert_called_once_with(
mock_predicted_epc,
property_id=PROPERTY_ID,
portfolio_id=PORTFOLIO_ID,
source="predicted",
)
mock_uow.plan.save.assert_called_once()
mock_uow.commit.assert_called_once()
# ---------------------------------------------------------------------------
# Partial batch failure — per-property isolation
# ---------------------------------------------------------------------------
def test_per_property_failure_fails_child_subtask_and_siblings_continue() -> None:
"""Two properties: property 1 succeeds, property 2 fails during modelling.
The handler does not raise; both properties are run in the one batched
run_subtasks pass, and property 1's write is still committed."""
# Arrange
pid1, pid2 = 111, 222
mock_engine = _engine_mock([pid1, pid2], [1001, 1002], [POSTCODE, POSTCODE])
mock_plan = _plan_mock()
mock_uow = MagicMock()
call_count = {"n": 0}
def _run_modelling_side_effect(*args: Any, **kwargs: Any) -> Any:
call_count["n"] += 1
if call_count["n"] == 2:
raise ValueError("modelling exploded")
return mock_plan
mock_orch = _mock_orchestrator()
with ExitStack() as stack:
stack.enter_context(patch("applications.modelling_e2e.handler.os.environ", _ENV))
stack.enter_context(
patch("applications.modelling_e2e.handler._get_engine", return_value=mock_engine)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(patch("applications.modelling_e2e.handler.GeospatialS3Repository"))
stack.enter_context(patch("applications.modelling_e2e.handler.GoogleSolarApiClient"))
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler._solar_insights_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"))
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
side_effect=_run_modelling_side_effect,
)
)
MockUoW = stack.enter_context(patch("applications.modelling_e2e.handler.PostgresUnitOfWork"))
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act — must not raise even though pid2 fails
_call_handler(
{"property_ids": [pid1, pid2], "portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID, "no_solar": True, "dry_run": False},
orchestrator=mock_orch,
)
# Both properties run in one batched run_subtasks pass; pid1 still committed
# (the failed pid2 is isolated and simply never queued for write).
mock_orch.run_subtasks.assert_called_once()
assert len(mock_orch.run_subtasks.call_args.args[1]) == 2
mock_uow.commit.assert_called_once()
# ---------------------------------------------------------------------------
# End-of-batch write — one transaction for the whole batch
# ---------------------------------------------------------------------------
def test_batch_persists_in_one_transaction_and_one_baseline_run(
_baseline_orchestrator: MagicMock,
) -> None:
"""Three properties model independently but persist together: every Plan is
saved, yet the write Unit of Work commits exactly once and the Baseline
orchestrator runs once for the whole batch (the RDS-CPU win)."""
# Arrange
pid1, pid2, pid3 = 111, 222, 333
mock_engine = _engine_mock(
[pid1, pid2, pid3], [1001, 1002, 1003], [POSTCODE, POSTCODE, POSTCODE]
)
mock_uow = MagicMock()
with ExitStack() as stack:
stack.enter_context(patch("applications.modelling_e2e.handler.os.environ", _ENV))
stack.enter_context(
patch("applications.modelling_e2e.handler._get_engine", return_value=mock_engine)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(patch("applications.modelling_e2e.handler.GeospatialS3Repository"))
stack.enter_context(patch("applications.modelling_e2e.handler.GoogleSolarApiClient"))
stack.enter_context(
patch("applications.modelling_e2e.handler._spatial_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler._solar_insights_for", return_value=None)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.overlays_from", return_value=[])
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides"))
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch("applications.modelling_e2e.handler.run_modelling", return_value=_plan_mock())
)
MockUoW = stack.enter_context(patch("applications.modelling_e2e.handler.PostgresUnitOfWork"))
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
_call_handler(
{"property_ids": [pid1, pid2, pid3], "portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID, "no_solar": True, "dry_run": False}
)
# Assert — all three Plans saved, but a single shared transaction:
assert mock_uow.plan.save.call_count == 3
assert mock_uow.property.mark_modelled.call_count == 3
mock_uow.commit.assert_called_once()
# One write Unit of Work opened for the whole batch, not one per property.
MockUoW.return_value.__enter__.assert_called_once()
# Baseline re-established once, for every written property together.
_baseline_orchestrator.return_value.run.assert_called_once_with([pid1, pid2, pid3])
# ---------------------------------------------------------------------------
# Cohort cache hit
# ---------------------------------------------------------------------------
def test_cohort_cache_prevents_duplicate_candidates_for_calls() -> None:
"""Two properties sharing a postcode: candidates_for is fetched once and
cached — EpcComparablePropertiesRepository.candidates_for called once."""
# Arrange
pid1, pid2 = 301, 302
mock_engine = _engine_mock(
[pid1, pid2], [3001, 3002], [POSTCODE, POSTCODE]
)
mock_plan = _plan_mock()
mock_uow = MagicMock()
mock_predicted_epc = MagicMock()
from datatypes.epc.domain.epc_property_data import BuildingPartIdentifier
mock_part = MagicMock()
mock_part.identifier = BuildingPartIdentifier.MAIN
mock_predicted_epc.sap_building_parts = [mock_part]
mock_comparables = MagicMock()
mock_comparables.members = [MagicMock()]
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = None # force prediction path
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._spatial_for", return_value=None
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.overlays_from", return_value=[]
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
from domain.epc_prediction.prediction_target import PredictionTargetAttributes
stack.enter_context(
patch(
"applications.modelling_e2e.handler.OverrideBackedPredictionAttributesReader"
)
).return_value.attributes_for.return_value = PredictionTargetAttributes(
property_type="2"
)
stack.enter_context(
patch("applications.modelling_e2e.handler.select_comparables")
).return_value = mock_comparables
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcPrediction")
).return_value.predict.return_value = mock_predicted_epc
MockCandidates = stack.enter_context(
patch(
"applications.modelling_e2e.handler.EpcComparablePropertiesRepository"
)
)
MockCandidates.return_value.candidates_for.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides")
)
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=mock_plan,
)
)
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
MockUoW.return_value.__enter__.return_value = mock_uow
MockUoW.return_value.__exit__.return_value = False
# Act
_call_handler(
{
"property_ids": [pid1, pid2],
"portfolio_id": PORTFOLIO_ID,
"scenario_id": SCENARIO_ID,
"no_solar": True,
"dry_run": False,
}
)
# Assert — cohort fetched exactly once despite two properties
MockCandidates.return_value.candidates_for.assert_called_once_with(POSTCODE)
# ---------------------------------------------------------------------------
# Dry-run
# ---------------------------------------------------------------------------
def test_dry_run_skips_all_db_writes() -> None:
"""dry_run=True: run_modelling executes but PostgresUnitOfWork is never
entered — no DB writes occur for any property in the batch."""
# Arrange
mock_engine = _engine_mock([PROPERTY_ID], [UPRN], [POSTCODE])
with ExitStack() as stack:
stack.enter_context(
patch("applications.modelling_e2e.handler.os.environ", _ENV)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._get_engine",
return_value=mock_engine,
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.EpcClientService")
).return_value.get_by_uprn.return_value = MagicMock()
stack.enter_context(
patch("applications.modelling_e2e.handler.GeospatialS3Repository")
)
stack.enter_context(
patch("applications.modelling_e2e.handler.GoogleSolarApiClient")
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._spatial_for", return_value=None
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler._solar_insights_for",
return_value=None,
)
)
stack.enter_context(
patch(
"applications.modelling_e2e.handler.overlays_from", return_value=[]
)
)
stack.enter_context(
patch("applications.modelling_e2e.handler.PropertyOverridesPostgresReader")
).return_value.overrides_for_many.return_value = {}
stack.enter_context(
patch("applications.modelling_e2e.handler.ScenarioPostgresRepository")
).return_value.get_many.return_value = [MagicMock()]
stack.enter_context(
patch("applications.modelling_e2e.handler.catalogue_snapshot_with_off_catalogue_overrides")
)
stack.enter_context(patch("applications.modelling_e2e.handler.Session"))
stack.enter_context(
patch(
"applications.modelling_e2e.handler.run_modelling",
return_value=_plan_mock(),
)
)
MockUoW = stack.enter_context(
patch("applications.modelling_e2e.handler.PostgresUnitOfWork")
)
# Act
_call_handler({**_BODY, "dry_run": True})
# Assert — UoW never entered
MockUoW.return_value.__enter__.assert_not_called()
Reference in a new issue View git blame Copy permalink