Update tests to work with new formula engine interface

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

Author: shsms

examples/power_distribution.py

@@ -230,7 +230,7 @@ async def run() -> None:
 
     client_actor = DataCollectingActor(
         request_channel=power_dist_req_chan.new_sender(),
-        active_power_data=await logical_meter.grid_power(),
+        active_power_data=logical_meter.grid_power().new_receiver(),
     )
 
     await actor.run(service_actor, client_actor, power_distributor)

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -29,15 +29,17 @@ async def test_formula_composition(  # pylint: disable=too-many-locals
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        main_meter_recv = await get_resampled_stream(
+        main_meter_recv = get_resampled_stream(
             4,
             ComponentMetricId.ACTIVE_POWER,
         )
-        grid_power_recv = await logical_meter.grid_power()
-        battery_power_recv = await logical_meter.battery_power()
-        pv_power_recv = await logical_meter.pv_power()
+        grid_power_recv = logical_meter.grid_power().new_receiver()
+        battery_power_recv = logical_meter.battery_power().new_receiver()
+        pv_power_recv = logical_meter.pv_power().new_receiver()
 
-        engine = (pv_power_recv.clone() + battery_power_recv.clone()).build("inv_power")
+        engine = (logical_meter.pv_power() + logical_meter.battery_power()).build(
+            "inv_power"
+        )
         inv_calc_recv = engine.new_receiver()
 
         count = 0
@@ -69,9 +71,11 @@ async def test_formula_composition_missing_pv(self, mocker: MockerFixture) -> No
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        battery_power_recv = await logical_meter.battery_power()
-        pv_power_recv = await logical_meter.pv_power()
-        engine = (pv_power_recv.clone() + battery_power_recv.clone()).build("inv_power")
+        battery_power_recv = logical_meter.battery_power().new_receiver()
+        pv_power_recv = logical_meter.pv_power().new_receiver()
+        engine = (logical_meter.pv_power() + logical_meter.battery_power()).build(
+            "inv_power"
+        )
         inv_calc_recv = engine.new_receiver()
 
         count = 0
@@ -96,9 +100,11 @@ async def test_formula_composition_missing_bat(self, mocker: MockerFixture) -> N
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        battery_power_recv = await logical_meter.battery_power()
-        pv_power_recv = await logical_meter.pv_power()
-        engine = (pv_power_recv.clone() + battery_power_recv.clone()).build("inv_power")
+        battery_power_recv = logical_meter.battery_power().new_receiver()
+        pv_power_recv = logical_meter.pv_power().new_receiver()
+        engine = (logical_meter.pv_power() + logical_meter.battery_power()).build(
+            "inv_power"
+        )
         inv_calc_recv = engine.new_receiver()
 
         count = 0
@@ -125,12 +131,10 @@ async def test_3_phase_formulas(self, mocker: MockerFixture) -> None:
         logical_meter = microgrid.logical_meter()
         ev_pool = microgrid.ev_charger_pool()
 
-        grid_current_recv = await logical_meter.grid_current()
-        ev_current_recv = await ev_pool.current()
+        grid_current_recv = logical_meter.grid_current().new_receiver()
+        ev_current_recv = ev_pool.current().new_receiver()
 
-        engine = (grid_current_recv.clone() - ev_current_recv.clone()).build(
-            "net_current"
-        )
+        engine = (logical_meter.grid_current() - ev_pool.current()).build("net_current")
         net_current_recv = engine.new_receiver()
 
         await synchronize_receivers(

tests/timeseries/_formula_engine/utils.py

@@ -7,20 +7,17 @@
 
 from datetime import datetime
 from math import isclose
+from typing import overload
 
 from frequenz.channels import Receiver
 
 from frequenz.sdk.microgrid import _data_pipeline
 from frequenz.sdk.microgrid.component import ComponentMetricId
 from frequenz.sdk.timeseries import Sample, Sample3Phase
-from frequenz.sdk.timeseries._formula_engine import (
-    FormulaReceiver,
-    FormulaReceiver3Phase,
-    ResampledFormulaBuilder,
-)
+from frequenz.sdk.timeseries._formula_engine import ResampledFormulaBuilder
 
 
-async def get_resampled_stream(  # pylint: disable=too-many-arguments
+def get_resampled_stream(  # pylint: disable=too-many-arguments
     comp_id: int,
     metric_id: ComponentMetricId,
 ) -> Receiver[Sample]:
@@ -36,7 +33,7 @@ async def get_resampled_stream(  # pylint: disable=too-many-arguments
         _data_pipeline._get()._resampling_request_sender(),
         metric_id,
     )
-    return await builder._get_resampled_receiver(
+    return builder._get_resampled_receiver(
         comp_id,
         metric_id,
     )
@@ -52,16 +49,28 @@ def equal_float_lists(list1: list[float], list2: list[float]) -> bool:
     )
 
 
+@overload
 async def synchronize_receivers(
-    receivers: list[FormulaReceiver | FormulaReceiver3Phase | Receiver[Sample]],
+    receivers: list[Receiver[Sample]],
+) -> None:
+    ...
+
+
+@overload
+async def synchronize_receivers(
+    receivers: list[Receiver[Sample3Phase]],
+) -> None:
+    ...
+
+
+async def synchronize_receivers(
+    receivers: list[Receiver[Sample]] | list[Receiver[Sample3Phase]],
 ) -> None:
     """Check if given receivers are all returning the same timestamp.
 
     If not, try to synchronize them.
     """
-    by_ts: dict[
-        datetime, list[FormulaReceiver | FormulaReceiver3Phase | Receiver[Sample]]
-    ] = {}
+    by_ts: dict[datetime, list[Receiver[Sample] | Receiver[Sample3Phase]]] = {}
     for recv in receivers:
         while True:
             sample = await recv.receive()

tests/timeseries/test_ev_charger_pool.py

@@ -92,12 +92,12 @@ async def test_ev_power(  # pylint: disable=too-many-locals
 
         ev_pool = microgrid.ev_charger_pool()
 
-        main_meter_recv = await get_resampled_stream(
+        main_meter_recv = get_resampled_stream(
             mockgrid.main_meter_id,
             ComponentMetricId.ACTIVE_POWER,
         )
-        grid_power_recv = await logical_meter.grid_power()
-        ev_power_recv = await ev_pool.power()
+        grid_power_recv = logical_meter.grid_power().new_receiver()
+        ev_power_recv = ev_pool.power().new_receiver()
 
         await synchronize_receivers([grid_power_recv, main_meter_recv, ev_power_recv])
 
@@ -158,7 +158,7 @@ async def mock_current_streams(
             mock_current_streams,
         )
 
-        recv = await ev_pool.component_data(evc_id)
+        recv = ev_pool.component_data(evc_id)
 
         await send_resampled_current(2, 3, 5)
         await asyncio.sleep(0.02)

tests/timeseries/test_formula_engine.py

@@ -13,7 +13,6 @@
 from frequenz.sdk.timeseries._formula_engine._formula_engine import (
     FormulaBuilder,
     FormulaEngine,
-    FormulaReceiver,
     HigherOrderFormulaBuilder,
 )
 from frequenz.sdk.timeseries._formula_engine._tokenizer import (
@@ -48,7 +47,7 @@ def test_1(self) -> None:
 class TestFormulaEngine:
     """Tests for the FormulaEngine."""
 
-    async def run_test(
+    async def run_test(  # pylint: disable=too-many-locals
         self,
         formula: str,
         postfix: str,
@@ -70,11 +69,9 @@ async def run_test(
             elif token.type == TokenType.OPER:
                 builder.push_oper(token.value)
         engine = builder.build()
+        results_rx = engine.new_receiver()
 
-        assert (
-            repr(engine._evaluator._steps)  # pylint: disable=protected-access
-            == postfix
-        )
+        assert repr(builder._steps) == postfix  # pylint: disable=protected-access
 
         now = datetime.now()
         tests_passed = 0
@@ -86,9 +83,7 @@ async def run_test(
                     for chan, value in zip(channels.values(), io_input)
                 ]
             )
-            next_val = (
-                await engine._evaluator.apply()  # pylint: disable=protected-access
-            )
+            next_val = await results_rx.receive()
             assert (next_val).value == io_output
             tests_passed += 1
         await engine._stop()  # pylint: disable=protected-access
@@ -293,7 +288,7 @@ async def test_nones_are_not_zeros(self) -> None:
         )
 
 
-class TestFormulaChannel:
+class TestFormulaEngineComposition:
     """Tests for formula channels."""
 
     def make_engine(self, stream_id: int, data: Receiver[Sample]) -> FormulaEngine:
@@ -305,18 +300,18 @@ def make_engine(self, stream_id: int, data: Receiver[Sample]) -> FormulaEngine:
             data,
             nones_are_zeros=False,
         )
-        return builder.build()
+        return FormulaEngine(builder)
 
     async def run_test(  # pylint: disable=too-many-locals
         self,
         num_items: int,
         make_builder: Union[
             Callable[
-                [FormulaReceiver, FormulaReceiver, FormulaReceiver],
+                [FormulaEngine, FormulaEngine, FormulaEngine],
                 HigherOrderFormulaBuilder,
             ],
             Callable[
-                [FormulaReceiver, FormulaReceiver, FormulaReceiver, FormulaReceiver],
+                [FormulaEngine, FormulaEngine, FormulaEngine, FormulaEngine],
                 HigherOrderFormulaBuilder,
             ],
         ],
@@ -329,7 +324,7 @@ async def run_test(  # pylint: disable=too-many-locals
             self.make_engine(ctr, channels[ctr].new_receiver())
             for ctr in range(num_items)
         ]
-        builder = make_builder(*[e.new_receiver() for e in l1_engines])
+        builder = make_builder(*l1_engines)
         engine = builder.build("l2 formula", nones_are_zeros)
         result_chan = engine.new_receiver()
 
@@ -567,7 +562,7 @@ async def run_test(
 
         builder.push_average(streams)
         engine = builder.build()
-
+        results_rx = engine.new_receiver()
         now = datetime.now()
         tests_passed = 0
         for io_pair in io_pairs:
@@ -578,9 +573,7 @@ async def run_test(
                     for chan, value in zip(channels.values(), io_input)
                 ]
             )
-            next_val = (
-                await engine._evaluator.apply()  # pylint: disable=protected-access
-            )
+            next_val = await results_rx.receive()
             assert (next_val).value == io_output
             tests_passed += 1
         await engine._stop()  # pylint: disable=protected-access

tests/timeseries/test_logical_meter.py

@@ -31,9 +31,9 @@ async def test_grid_power_1(self, mocker: MockerFixture) -> None:
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        grid_power_recv = await logical_meter.grid_power()
+        grid_power_recv = logical_meter.grid_power().new_receiver()
 
-        main_meter_recv = await get_resampled_stream(
+        main_meter_recv = get_resampled_stream(
             mockgrid.main_meter_id,
             ComponentMetricId.ACTIVE_POWER,
         )
@@ -63,10 +63,10 @@ async def test_grid_power_2(
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        grid_power_recv = await logical_meter.grid_power()
+        grid_power_recv = logical_meter.grid_power().new_receiver()
 
         meter_receivers = [
-            await get_resampled_stream(
+            get_resampled_stream(
                 meter_id,
                 ComponentMetricId.ACTIVE_POWER,
             )
@@ -105,19 +105,19 @@ async def test_battery_and_pv_power(  # pylint: disable=too-many-locals
         await mockgrid.start(mocker)
         logical_meter = microgrid.logical_meter()
 
-        battery_power_recv = await logical_meter.battery_power()
-        pv_power_recv = await logical_meter.pv_power()
+        battery_power_recv = logical_meter.battery_power().new_receiver()
+        pv_power_recv = logical_meter.pv_power().new_receiver()
 
         bat_inv_receivers = [
-            await get_resampled_stream(
+            get_resampled_stream(
                 meter_id,
                 ComponentMetricId.ACTIVE_POWER,
             )
             for meter_id in mockgrid.battery_inverter_ids
         ]
 
         pv_inv_receivers = [
-            await get_resampled_stream(
+            get_resampled_stream(
                 meter_id,
                 ComponentMetricId.ACTIVE_POWER,
             )