refactor processor

backend/app/db/functions/recommendations_functions.py

@@ -621,7 +621,7 @@ def get_plans_by_portfolio_id(portfolio_id: int) -> List[PlanModel]:
     raise NotImplementedError
 
 
-def get_scenario(scenario_id: int) -> List[ScenarioModel]:
+def get_scenario(scenario_id: int) -> ScenarioModel:
     raise NotImplementedError
 
 

backend/app/domain/classes/plan.py

@@ -11,12 +11,15 @@ class Plan:
     def __init__(
         self, record: PlanRecord, scenario: Scenario, id: Optional[int] = None
     ):
-        self.id = id
-        self._record = record
-        self.scenario = scenario
+        self.id: Optional[int] = id
+        self.record: PlanRecord = record
+        self.scenario: Scenario = scenario
 
     @classmethod
     def from_sqlalchemy(cls, plan_model: PlanModel, scenario: Scenario) -> Plan:
+        if not scenario:
+            raise ValueError(f"No Scenario associated with Plan of ID {plan_model.id}")
+
         record = PlanRecord(
             property_id=plan_model.property_id,
             portfolio_id=plan_model.portfolio_id,
@@ -43,4 +46,4 @@ class Plan:
         return cls(record=record, scenario=scenario, id=plan_model.id)
 
     def set_default(self, value: bool) -> None:
-        self._record = replace(self._record, is_default=value)
+        self.record = replace(self.record, is_default=value)

backend/categorisation/categorisation_logic.py

@@ -1,12 +1,12 @@
 from typing import List
-from backend.app.db.models.recommendations import PlanModel
+from backend.app.domain.classes.plan import Plan
 
 
 class CategorisationLogic:
     @staticmethod
-    def get_compliant_plans(plans: List[PlanModel]) -> List[PlanModel]:
+    def get_compliant_plans(plans: List[Plan]) -> List[Plan]:
         raise NotImplementedError
 
     @staticmethod
-    def get_cheapest_plan(plans: List[PlanModel]) -> PlanModel:
+    def get_cheapest_plan(plans: List[Plan]) -> Plan:
         raise NotImplementedError

backend/categorisation/processor.py

@@ -1,35 +1,64 @@
+from collections import defaultdict
 from typing import List
 
 from backend.app.db.functions.recommendations_functions import (
     get_plans_by_portfolio_id,
-    get_property_ids,
+    get_scenario,
     set_plan_default,
 )
-from backend.app.db.models.recommendations import PlanModel
+from backend.app.domain.classes.plan import Plan
 from backend.categorisation.categorisation_logic import CategorisationLogic
+from utils.logger import setup_logger
+
+logger = setup_logger()
 
 
 def process_portfolio(portfolio_id: int) -> None:
-    # Get all plans (including scenarios) for all properties in the portfolio
-    plans: List[PlanModel] = get_plans_by_portfolio_id(portfolio_id)
+    plans = _load_plans_for_portfolio(portfolio_id)
+    plans_by_property = _group_plans_by_property(plans)
 
-    # For each property, get all compliant plans
-    property_ids: List[int] = get_property_ids(portfolio_id)
+    for property_plans in plans_by_property.values():
+        cheapest_plan = _choose_cheapest_relevant_plan(property_plans)
+        _update_default_flags(property_plans, cheapest_plan)
 
-    # For each property, find the cheapest compliant plan
-    for id in property_ids:
-        plans_for_property: List[PlanModel] = [
-            plan for plan in plans if plan.property_id == id
-        ]
 
-        compliant_plans_for_property: List[PlanModel] = (
-            CategorisationLogic.get_compliant_plans(plans_for_property)
+def _load_plans_for_portfolio(portfolio_id: int) -> List[Plan]:
+    plan_models = get_plans_by_portfolio_id(portfolio_id)
+    plans: List[Plan] = []
+
+    for model in plan_models:
+        if not model.scenario_id:
+            logger.info(f"No Scenario associated with Plan of ID {model.id}")
+            continue
+
+        scenario_model = get_scenario(model.scenario_id)
+        plans.append(Plan.from_sqlalchemy(model, scenario_model))
+
+    return plans
+
+
+def _group_plans_by_property(plans: List[Plan]) -> dict[int, List[Plan]]:
+    grouped: dict[int, List[Plan]] = defaultdict(list)
+
+    for plan in plans:
+        grouped[plan.record.property_id].append(plan)
+
+    return grouped
+
+
+def _choose_cheapest_relevant_plan(plans: List[Plan]) -> Plan:
+    compliant_plans = CategorisationLogic.get_compliant_plans(plans)
+
+    plans_to_consider = compliant_plans or plans
+    return CategorisationLogic.get_cheapest_plan(plans_to_consider)
+
+
+def _update_default_flags(plans: List[Plan], cheapest_plan: Plan) -> None:
+    for plan in plans:
+        if plan.id is None:
+            raise ValueError("Cannot update Plan with missing ID")
+
+        set_plan_default(
+            plan.id,
+            plan.id == cheapest_plan.id,
         )
-
-        # Choose cheapest compliant plan, or fallback to cheapest overall plan
-        plans_to_consider = compliant_plans_for_property or plans_for_property
-        cheapest_plan = CategorisationLogic.get_cheapest_plan(plans_to_consider)
-
-        # Update DB: set is_default = True for cheapest plan, False for others
-        for plan in plans_for_property:
-            set_plan_default(plan.id, plan.id == cheapest_plan.id)