diff --git a/applications/modelling_e2e/handler.py b/applications/modelling_e2e/handler.py index 55a646676..51c3e8149 100644 --- a/applications/modelling_e2e/handler.py +++ b/applications/modelling_e2e/handler.py @@ -210,7 +210,7 @@ def _flush_writes(engine: Engine, writes: list[_PropertyWrite]) -> None: # p.is_new rule — readers select one default Plan per property, so a # first run under a non-default Scenario must still produce one). # Also heals properties left default-less by runs before this rule. - have_default: set[int] = uow.plan.property_ids_with_default_plans( + pids_with_default: set[int] = uow.plan.property_ids_with_default_plans( [w.property_id for w in writes] ) plan_requests = [ @@ -219,7 +219,7 @@ def _flush_writes(engine: Engine, writes: list[_PropertyWrite]) -> None: property_id=w.property_id, scenario_id=w.scenario_id, portfolio_id=w.portfolio_id, - is_default=w.is_default or w.property_id not in have_default, + is_default=w.is_default or w.property_id not in pids_with_default, ) for w in writes ] diff --git a/repositories/plan/plan_postgres_repository.py b/repositories/plan/plan_postgres_repository.py index 21540523a..dc3b135a7 100644 --- a/repositories/plan/plan_postgres_repository.py +++ b/repositories/plan/plan_postgres_repository.py @@ -47,6 +47,8 @@ class PlanPostgresRepository(PlanRepository): )[0] def property_ids_with_default_plans(self, property_ids: list[int]) -> set[int]: + if not property_ids: + return set() rows = self._session.exec( select(col(PlanModel.property_id)) .distinct() diff --git a/tests/orchestration/fakes.py b/tests/orchestration/fakes.py index 52f409693..4c6d323b9 100644 --- a/tests/orchestration/fakes.py +++ b/tests/orchestration/fakes.py @@ -212,6 +212,7 @@ class FakePlanRepository(PlanRepository): def __init__(self) -> None: self.saved: dict[tuple[int, int], Plan] = {} + self.default_property_ids: set[int] = set() self._next_id = 1 def save( @@ -224,6 +225,8 @@ class FakePlanRepository(PlanRepository): is_default: bool, ) -> int: self.saved[(property_id, scenario_id)] = plan + if is_default: + self.default_property_ids.add(property_id) plan_id = self._next_id self._next_id += 1 return plan_id @@ -241,9 +244,7 @@ class FakePlanRepository(PlanRepository): ] def property_ids_with_default_plans(self, property_ids: list[int]) -> set[int]: - # The fake store does not track defaults; every property reads as - # first-time (no default yet). - return set() + return self.default_property_ids.intersection(property_ids) class _UnsetProductRepo(ProductRepository):