diff --git a/applications/modelling_e2e/handler.py b/applications/modelling_e2e/handler.py index 9a8c8e73..30c26f82 100644 --- a/applications/modelling_e2e/handler.py +++ b/applications/modelling_e2e/handler.py @@ -309,6 +309,7 @@ def handler(body: dict[str, Any], context: Any) -> Optional[dict[str, Any]]: epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn) overrides = overlays_from(overrides_reader.overrides_for(property_id)) + predicted_epc: Optional[EpcPropertyData] = None if epc is not None: logger.info(f"property={property_id} lodged EPC found") @@ -391,6 +392,16 @@ def handler(body: dict[str, Any], context: Any) -> Optional[dict[str, Any]]: uow.epc.save( epc, property_id=property_id, portfolio_id=portfolio_id ) + elif predicted_epc is not None: + # Persist the synthesised EPC in the predicted slot (ADR-0031), + # so the Baseline stage can re-hydrate it and downstream sees + # the picture the Plan was modelled from. + uow.epc.save( + predicted_epc, + property_id=property_id, + portfolio_id=portfolio_id, + source="predicted", + ) if spatial is not None: uow.spatial.save(uprn, spatial) if ( @@ -417,12 +428,12 @@ def handler(body: dict[str, Any], context: Any) -> Optional[dict[str, Any]]: uow.commit() logger.info(f"property={property_id} plan saved") - # Baseline Performance is re-established from the persisted EPC, so - # it runs after the Plan UoW commits. Only lodged Properties have - # the lodged figures the Baseline reads; predicted ones are skipped. - if epc is not None: - baseline_orchestrator.run([property_id]) - logger.info(f"property={property_id} baseline saved") + # Baseline Performance is re-established from the persisted EPC + # (lodged or predicted), so it runs after the Plan UoW commits. By + # here the property always has a persisted EPC — a property that + # could be neither fetched nor predicted raised earlier. + baseline_orchestrator.run([property_id]) + logger.info(f"property={property_id} baseline saved") except Exception as error: # noqa: BLE001 logger.error( diff --git a/tests/applications/modelling_e2e/test_handler.py b/tests/applications/modelling_e2e/test_handler.py index 6929d13f..6478572a 100644 --- a/tests/applications/modelling_e2e/test_handler.py +++ b/tests/applications/modelling_e2e/test_handler.py @@ -312,12 +312,12 @@ def test_skipped_cohort_certs_are_surfaced_in_the_outputs() -> None: # --------------------------------------------------------------------------- -def test_prediction_path_saves_plan_without_epc_save( +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 and saves the plan — but never calls epc.save, and (having no - lodged figures) never establishes a Baseline.""" + 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() @@ -417,11 +417,16 @@ def test_prediction_path_saves_plan_without_epc_save( # Act _call_handler(_BODY) - # Assert — epc.save NOT called (no lodged cert), plan IS saved, no baseline - mock_uow.epc.save.assert_not_called() + # 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_not_called() + _baseline_orchestrator.return_value.run.assert_called_once_with([PROPERTY_ID]) # --------------------------------------------------------------------------- @@ -620,9 +625,15 @@ def test_empty_own_postcode_broadens_to_nearby_and_predicts() -> None: # Act _call_handler(_BODY) - # Assert — broadening fired, and the broadened cohort produced a saved plan. + # 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_not_called() # predicted, never lodged + 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()