Put frequently used code in base function.

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

src/frequenz/sdk/timeseries/_formula_engine/_formula_generators/_consumer_power_formula.py

@@ -37,20 +37,9 @@ def generate(self) -> FormulaEngine[Power]:
         builder = self._get_builder(
             "consumer-power", ComponentMetricId.ACTIVE_POWER, Power.from_watts
         )
-        component_graph = connection_manager.get().component_graph
-        grid_component = next(
-            (
-                comp
-                for comp in component_graph.components()
-                if comp.category == ComponentCategory.GRID
-            ),
-            None,
-        )
 
-        if grid_component is None:
-            raise ComponentNotFound("Grid component not found in the component graph.")
+        grid_successors = self._get_grid_component_successors()
 
-        grid_successors = component_graph.successors(grid_component.component_id)
         if not grid_successors:
             raise ComponentNotFound("No components found in the component graph.")
 

src/frequenz/sdk/timeseries/_formula_engine/_formula_generators/_formula_generator.py

@@ -15,6 +15,7 @@
 from frequenz.channels import Sender
 
 from ....actor import ChannelRegistry, ComponentMetricRequest
+from ....microgrid import component, connection_manager
 from ....microgrid.component import ComponentMetricId
 from ..._quantities import QuantityT
 from .._formula_engine import FormulaEngine, FormulaEngine3Phase
@@ -100,6 +101,36 @@ def _get_builder(
         )
         return builder
 
+    def _get_grid_component_successors(self) -> set[component.Component]:
+        """Get the set of grid component successors in the component graph.
+
+        Returns:
+            A set of grid component successors.
+
+        Raises:
+            ComponentNotFound: If the grid component is not found in the component graph.
+            ComponentNotFound: If no successor components are found in the component graph.
+        """
+        component_graph = connection_manager.get().component_graph
+        grid_component = next(
+            iter(
+                component_graph.components(
+                    component_category={component.ComponentCategory.GRID}
+                )
+            ),
+            None,
+        )
+
+        if grid_component is None:
+            raise ComponentNotFound("Grid component not found in the component graph.")
+
+        grid_successors = component_graph.successors(grid_component.component_id)
+
+        if not grid_successors:
+            raise ComponentNotFound("No components found in the component graph.")
+
+        return grid_successors
+
     @abstractmethod
     def generate(
         self,

src/frequenz/sdk/timeseries/_formula_engine/_formula_generators/_grid_current_formula.py

@@ -5,11 +5,10 @@
 
 from typing import Set
 
-from ....microgrid import connection_manager
 from ....microgrid.component import Component, ComponentCategory, ComponentMetricId
 from ..._quantities import Current
 from .._formula_engine import FormulaEngine, FormulaEngine3Phase
-from ._formula_generator import ComponentNotFound, FormulaGenerator
+from ._formula_generator import FormulaGenerator
 
 
 class GridCurrentFormula(FormulaGenerator[Current]):
@@ -24,22 +23,7 @@ def generate(self) -> FormulaEngine3Phase[Current]:
         Raises:
             ComponentNotFound: when the component graph doesn't have a `GRID` component.
         """
-        component_graph = connection_manager.get().component_graph
-        grid_component = next(
-            (
-                comp
-                for comp in component_graph.components()
-                if comp.category == ComponentCategory.GRID
-            ),
-            None,
-        )
-
-        if grid_component is None:
-            raise ComponentNotFound(
-                "Unable to find a GRID component from the component graph."
-            )
-
-        grid_successors = component_graph.successors(grid_component.component_id)
+        grid_successors = self._get_grid_component_successors()
 
         return FormulaEngine3Phase(
             "grid-current",

src/frequenz/sdk/timeseries/_formula_engine/_formula_generators/_grid_power_formula.py

@@ -3,11 +3,10 @@
 
 """Formula generator from component graph for Grid Power."""
 
-from ....microgrid import connection_manager
 from ....microgrid.component import ComponentCategory, ComponentMetricId
 from ..._quantities import Power
 from .._formula_engine import FormulaEngine
-from ._formula_generator import ComponentNotFound, FormulaGenerator, FormulaType
+from ._formula_generator import FormulaGenerator, FormulaType
 
 
 class GridPowerFormula(FormulaGenerator[Power]):
@@ -27,22 +26,7 @@ def generate(
         builder = self._get_builder(
             "grid-power", ComponentMetricId.ACTIVE_POWER, Power.from_watts
         )
-        component_graph = connection_manager.get().component_graph
-        grid_component = next(
-            (
-                comp
-                for comp in component_graph.components()
-                if comp.category == ComponentCategory.GRID
-            ),
-            None,
-        )
-
-        if grid_component is None:
-            raise ComponentNotFound(
-                "Unable to find a GRID component from the component graph."
-            )
-
-        grid_successors = component_graph.successors(grid_component.component_id)
+        grid_successors = self._get_grid_component_successors()
 
         # generate a formula that just adds values from all commponents that are
         # directly connected to the grid.  If the requested formula type is

src/frequenz/sdk/timeseries/_formula_engine/_formula_generators/_producer_power_formula.py

@@ -9,11 +9,7 @@
 from ....microgrid.component import ComponentCategory, ComponentMetricId
 from ..._quantities import Power
 from .._formula_engine import FormulaEngine
-from ._formula_generator import (
-    NON_EXISTING_COMPONENT_ID,
-    ComponentNotFound,
-    FormulaGenerator,
-)
+from ._formula_generator import NON_EXISTING_COMPONENT_ID, FormulaGenerator
 
 
 class ProducerPowerFormula(FormulaGenerator[Power]):
@@ -39,19 +35,7 @@ def generate(self) -> FormulaEngine[Power]:
             "producer_power", ComponentMetricId.ACTIVE_POWER, Power.from_watts
         )
         component_graph = connection_manager.get().component_graph
-        grid_component = next(
-            iter(
-                component_graph.components(component_category={ComponentCategory.GRID})
-            ),
-            None,
-        )
-
-        if grid_component is None:
-            raise ComponentNotFound("Grid component not found in the component graph.")
-
-        grid_successors = component_graph.successors(grid_component.component_id)
-        if not grid_successors:
-            raise ComponentNotFound("No components found in the component graph.")
+        grid_successors = self._get_grid_component_successors()
 
         if len(grid_successors) == 1:
             grid_meter = next(iter(grid_successors))