Fix bug where default plans were not being unset if they weren't included in scenarios to be considered

backend/app/db/functions/recommendations_functions.py

@@ -639,6 +639,13 @@ def get_plans_by_scenario_ids(ids: List[int]) -> List[PlanModel]:
         return session_any.exec(stmt).scalars().all()
 
 
+def get_plan_ids_by_scenario_ids(scenario_ids: List[int]) -> List[int]:
+    stmt = select(PlanModel.id).where(PlanModel.scenario_id.in_(scenario_ids))
+    with db_read_session() as session:
+        session_any: Any = session  # Typehint as Any to satisfy Pylance...
+        return session_any.exec(stmt).scalars().all()
+
+
 def get_scenarios_by_portfolio_id(portfolio_id: int) -> List[ScenarioModel]:
     stmt = select(ScenarioModel).where(ScenarioModel.portfolio_id == portfolio_id)
     with db_read_session() as session:
@@ -657,7 +664,19 @@ def get_default_plan_ids_for_property(property_id: int) -> List[int]:
     # This should in reality always return exactly 1 ID, but there's currently
     # no database constraint to enforce that, so account for 0 or >1
     stmt = select(PlanModel.id).where(
-        PlanModel.property_id == property_id and PlanModel.is_default
+        (PlanModel.property_id == property_id) & (PlanModel.is_default == True)
+    )
+    with db_read_session() as session:
+        session_any: Any = session  # Typehint as Any to satisfy Pylance...
+        return session_any.exec(stmt).scalars().all()
+
+
+def get_default_scenario_ids_for_portfolio(portfolio_id: int) -> List[int]:
+    # This should in reality always return exactly 1 ID, but there's currently
+    # no database constraint to enforce that, so account for 0 or >1
+    stmt = select(ScenarioModel.id).where(
+        (ScenarioModel.portfolio_id == portfolio_id)
+        & (ScenarioModel.is_default == True)
     )
     with db_read_session() as session:
         session_any: Any = session  # Typehint as Any to satisfy Pylance...

backend/categorisation/local_handler/invoke_local_lambda.py

@@ -10,7 +10,7 @@ payload = {
             "body": json.dumps(
                 {
                     "portfolio_id": 556,
-                    "scenarios_to_consider": [1040, 1041],
+                    "scenarios_to_consider": [],
                     "scenarios_priority_order": [],
                 }
             )

backend/categorisation/processor.py

@@ -3,7 +3,8 @@ from typing import Dict, List, Optional
 
 from backend.app.db.functions.recommendations_functions import (
     bulk_update_plans,
-    get_default_plan_ids_for_property,
+    get_default_scenario_ids_for_portfolio,
+    get_plan_ids_by_scenario_ids,
     get_plans_by_portfolio_id,
     get_plans_by_scenario_ids,
     get_scenarios_by_portfolio_id,
@@ -24,6 +25,17 @@ def process_portfolio(
 ) -> None:
     logger.info(f"Processing portfolio {portfolio_id}")
 
+    if scenarios_to_consider:
+        if len(scenarios_to_consider) < 2:
+            raise ValueError(
+                "Cannot run auto categorisation for fewer than 2 scenarios"
+            )
+
+    if scenarios_to_consider is not None:
+        _unset_defaults_for_scenarios_not_being_considered(
+            portfolio_id, scenarios_to_consider
+        )
+
     plans: List[Plan] = _load_plans_for_portfolio(portfolio_id, scenarios_to_consider)
 
     plans_by_property: Dict[int, List[Plan]] = _group_plans_by_property(plans)
@@ -37,11 +49,6 @@ def process_portfolio(
             property_plans, scenario_priority_order
         )
 
-        # Unset existing default(s) in case they are outside the plans to consider
-        default_plan_ids: List[int] = get_default_plan_ids_for_property(property_id)
-        for id in default_plan_ids:
-            set_plan_and_scenario_default(id, False)
-
         _update_default_flags(
             property_plans, cheapest_plan
         )  # TODO: we have already unset existing default(s), so this method can probably be a bit simpler now
@@ -87,16 +94,33 @@ def choose_cheapest_relevant_plan(
     return cheapest_plans[0]
 
 
+def _unset_defaults_for_scenarios_not_being_considered(
+    portfolio_id: int, scenarios_to_consider: List[int]
+) -> None:
+    default_scenario_ids: List[int] = get_default_scenario_ids_for_portfolio(
+        portfolio_id
+    )
+    scenarios_to_unset_default: List[int] = []
+
+    for id in default_scenario_ids:
+        if id not in scenarios_to_consider:
+            scenarios_to_unset_default.append(id)
+
+    logger.info(f"Scenarios to unset defaults: {scenarios_to_unset_default}")
+
+    if len(scenarios_to_unset_default) > 0:
+        plans_to_unset_default: List[int] = get_plan_ids_by_scenario_ids(
+            scenarios_to_unset_default
+        )
+        for plan_id in plans_to_unset_default:
+            set_plan_and_scenario_default(plan_id, False)  # TODO: do this in batch
+
+
 def _load_plans_for_portfolio(
     portfolio_id: int, scenarios_to_consider: Optional[List[int]] = None
 ) -> List[Plan]:
 
     if scenarios_to_consider:
-        if len(scenarios_to_consider) < 2:
-            raise ValueError(
-                "Cannot run auto categorisation for fewer than 2 scenarios"
-            )
-
         logger.info(f"Getting {len(scenarios_to_consider)} plans")
         plan_models: List[PlanModel] = get_plans_by_scenario_ids(scenarios_to_consider)