diff --git a/applications/modelling_e2e/handler.py b/applications/modelling_e2e/handler.py index c5aa189bf..829044a73 100644 --- a/applications/modelling_e2e/handler.py +++ b/applications/modelling_e2e/handler.py @@ -56,6 +56,7 @@ from domain.property_baseline.calculator_rebaseliner import CalculatorRebaseline from domain.sap10_calculator.calculator import Sap10Calculator from domain.tasks.tasks import Source from harness.console import run_modelling +from orchestration.task_orchestrator import TaskOrchestrator from orchestration.property_baseline_orchestrator import ( PropertyBaselineOrchestrator, ) @@ -97,6 +98,7 @@ from repositories.scenario.scenario_postgres_repository import ( ) from repositories.solar.solar_postgres_repository import SolarPostgresRepository from utilities.aws_lambda.task_handler import task_handler +from uuid import UUID from utilities.logger import setup_logger _engine: Optional[Engine] = None @@ -212,8 +214,8 @@ def _predict_epc( return predicted -@task_handler(task_source="modelling_e2e", source=Source.PROPERTY) -def handler(body: dict[str, Any], context: Any) -> Optional[dict[str, Any]]: +@task_handler(task_source="modelling_e2e", source=Source.PROPERTY, pass_task_orchestrator=True) +def handler(body: dict[str, Any], context: Any, orchestrator: TaskOrchestrator, task_id: UUID) -> None: trigger = ModellingE2ETriggerBody.model_validate(body) property_ids = trigger.property_ids portfolio_id = trigger.portfolio_id diff --git a/tests/applications/modelling_e2e/test_handler.py b/tests/applications/modelling_e2e/test_handler.py index baa201031..2ac4520d5 100644 --- a/tests/applications/modelling_e2e/test_handler.py +++ b/tests/applications/modelling_e2e/test_handler.py @@ -8,8 +8,9 @@ is needed. One test per distinct behaviour path. from __future__ import annotations from contextlib import ExitStack -from typing import Any, Iterator +from typing import Any, Iterator, Optional from unittest.mock import MagicMock, call, patch +from uuid import UUID, uuid4 import pytest from pydantic import ValidationError @@ -39,10 +40,22 @@ _BODY = { } -def _call_handler(body: dict[str, Any]) -> Any: +def _mock_orchestrator() -> MagicMock: + mock = MagicMock() + mock.run_subtask.side_effect = lambda subtask_id, work, **kwargs: work() + child = MagicMock() + child.id = uuid4() + mock.create_child_subtask.return_value = child + return mock + + +def _call_handler( + body: dict[str, Any], orchestrator: Optional[MagicMock] = None +) -> Any: from applications.modelling_e2e.handler import handler - return handler.__wrapped__(body, None) # type: ignore[attr-defined] + orch = orchestrator if orchestrator is not None else _mock_orchestrator() + return handler.__wrapped__(body, None, orch, uuid4()) # type: ignore[attr-defined] def _engine_mock( @@ -122,6 +135,72 @@ def test_trigger_body_rejects_missing_property_ids() -> None: ) +# --------------------------------------------------------------------------- +# 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")) + 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_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 — one child SubTask per property, inputs record the property_id + assert mock_orch.create_child_subtask.call_count == 3 + calls = mock_orch.create_child_subtask.call_args_list + recorded_ids = [c.kwargs["inputs"]["property_id"] for c in calls] + assert recorded_ids == [pid1, pid2, pid3] + # All three calls used the same task_id + assert all(c.args[0] == task_id for c in calls) + + # --------------------------------------------------------------------------- # Lodged EPC path # ---------------------------------------------------------------------------