Update FormulaEngine and LogicalMeter to use FormulaChannels

... instead of `Broadcast`s.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/logical_meter/_formula_engine.py

@@ -64,7 +64,7 @@ def __init__(
         self._steps = steps
         self._metric_fetchers = metric_fetchers
         self._first_run = True
-        self._channel = Broadcast[Sample](self._name)
+        self._channel = FormulaChannel(self._name, self)
         self._task = None
 
     async def _synchronize_metric_timestamps(
@@ -170,7 +170,9 @@ async def _run(self) -> None:
             else:
                 await sender.send(msg)
 
-    def new_receiver(self) -> Receiver[Sample]:
+    def new_receiver(
+        self, name: Optional[str] = None, max_size: int = 50
+    ) -> FormulaReceiver:
         """Create a new receiver that streams the output of the formula engine.
 
         Args:
@@ -183,7 +185,7 @@ def new_receiver(self) -> Receiver[Sample]:
         if self._task is None:
             self._task = asyncio.create_task(self._run())
 
-        return self._channel.new_receiver()
+        return self._channel.new_receiver(name, max_size)
 
 
 class FormulaBuilder:
@@ -585,3 +587,24 @@ def build(self, name: str, nones_are_zeros: bool = False) -> FormulaEngine:
         self._engine = builder.build()
 
         return self._engine
+
+    def new_receiver(
+        self, name: Optional[str] = None, max_size: int = 50
+    ) -> FormulaReceiver:
+        """Get a new receiver from the corresponding engine.
+
+        Args:
+            name: optional name for the receiver.
+            max_size: size of the receiver's buffer.
+
+        Returns:
+            A FormulaReceiver that streams formula output `Sample`s.
+
+        Raises:
+            RuntimeError: If `build` hasn't been called yet.
+        """
+        if self._engine is None:
+            raise RuntimeError(
+                "Please call `build()` first, before calls to `new_receiver()`"
+            )
+        return self._engine.new_receiver(name, max_size)

src/frequenz/sdk/timeseries/logical_meter/_logical_meter.py

@@ -10,13 +10,12 @@
 import uuid
 from typing import Dict, List, Type
 
-from frequenz.channels import Receiver, Sender
+from frequenz.channels import Sender
 
 from ...actor import ChannelRegistry, ComponentMetricRequest
 from ...microgrid import ComponentGraph
 from ...microgrid.component import ComponentMetricId
-from .. import Sample
-from ._formula_engine import FormulaEngine
+from ._formula_engine import FormulaEngine, FormulaReceiver
 from ._formula_generators import (
     BatteryPowerFormula,
     BatterySoCFormula,
@@ -83,7 +82,7 @@ async def start_formula(
         formula: str,
         component_metric_id: ComponentMetricId,
         nones_are_zeros: bool = False,
-    ) -> Receiver[Sample]:
+    ) -> FormulaReceiver:
         """Start execution of the given formula name.
 
         Args:
@@ -94,7 +93,7 @@ async def start_formula(
                 False, the returned value will be a None.
 
         Returns:
-            A Receiver that streams values with the formulas applied.
+            A FormulaReceiver that streams values with the formulas applied.
         """
         channel_key = formula + component_metric_id.value
         if channel_key in self._engines:
@@ -111,7 +110,7 @@ async def _get_formula_stream(
         self,
         channel_key: str,
         generator: Type[FormulaGenerator],
-    ) -> Receiver[Sample]:
+    ) -> FormulaReceiver:
         if channel_key in self._engines:
             return self._engines[channel_key].new_receiver()
 
@@ -121,7 +120,7 @@ async def _get_formula_stream(
         self._engines[channel_key] = engine
         return engine.new_receiver()
 
-    async def grid_power(self) -> Receiver[Sample]:
+    async def grid_power(self) -> FormulaReceiver:
         """Fetch the grid power for the microgrid.
 
         If a formula engine to calculate grid power is not already running, it
@@ -134,7 +133,7 @@ async def grid_power(self) -> Receiver[Sample]:
         """
         return await self._get_formula_stream("grid_power", GridPowerFormula)
 
-    async def battery_power(self) -> Receiver[Sample]:
+    async def battery_power(self) -> FormulaReceiver:
         """Fetch the cumulative battery power in the microgrid.
 
         If a formula engine to calculate cumulative battery power is not
@@ -147,7 +146,7 @@ async def battery_power(self) -> Receiver[Sample]:
         """
         return await self._get_formula_stream("battery_power", BatteryPowerFormula)
 
-    async def pv_power(self) -> Receiver[Sample]:
+    async def pv_power(self) -> FormulaReceiver:
         """Fetch the PV power production in the microgrid.
 
         If a formula engine to calculate PV power production is not
@@ -159,7 +158,7 @@ async def pv_power(self) -> Receiver[Sample]:
         """
         return await self._get_formula_stream("pv_power", PVPowerFormula)
 
-    async def _soc(self) -> Receiver[Sample]:
+    async def _soc(self) -> FormulaReceiver:
         """Fetch the SoC of the active batteries in the microgrid.
 
         NOTE: This method is part of the logical meter only temporarily, and will get

tests/timeseries/test_logical_meter.py

@@ -248,3 +248,55 @@ async def test_soc(self, mocker: MockerFixture) -> None:
             assert (await soc_recv.receive()).value == sum(bat_vals) / len(bat_vals)
 
         await mockgrid.cleanup()
+
+    async def test_formula_composition(  # pylint: disable=too-many-locals
+        self,
+        mocker: MockerFixture,
+    ) -> None:
+        """Test the battery power and pv power formulas."""
+        mockgrid = await MockMicrogrid.new(mocker, grid_side_meter=False)
+        mockgrid.add_batteries(3)
+        mockgrid.add_solar_inverters(2)
+        request_sender, channel_registry = await mockgrid.start()
+        logical_meter = LogicalMeter(
+            channel_registry,
+            request_sender,
+            microgrid.get().component_graph,
+        )
+
+        grid_power_recv = await logical_meter.grid_power()
+        battery_power_recv = await logical_meter.battery_power()
+        pv_power_recv = await logical_meter.pv_power()
+        main_meter_recv = await self._get_resampled_stream(
+            logical_meter,
+            channel_registry,
+            request_sender,
+            4,
+            ComponentMetricId.ACTIVE_POWER,
+        )
+
+        engine = (pv_power_recv.clone() + battery_power_recv.clone()).build("inv_power")
+        inv_calc_recv = engine.new_receiver()
+
+        count = 0
+        for _ in range(10):
+            grid_pow = await grid_power_recv.receive()
+            pv_pow = await pv_power_recv.receive()
+            bat_pow = await battery_power_recv.receive()
+            main_pow = await main_meter_recv.receive()
+            inv_calc_pow = await inv_calc_recv.receive()
+
+            assert grid_pow is not None and grid_pow.value is not None
+            assert inv_calc_pow is not None and inv_calc_pow.value is not None
+            assert bat_pow is not None and bat_pow.value is not None
+            assert pv_pow is not None and pv_pow.value is not None
+            assert main_pow is not None and main_pow.value is not None
+
+            assert inv_calc_pow.value == pv_pow.value + bat_pow.value
+            assert grid_pow.value == inv_calc_pow.value + main_pow.value
+            count += 1
+
+        await mockgrid.cleanup()
+        await engine._stop()  # pylint: disable=protected-access
+
+        assert count == 10