Mock Resampler in formula tests (frequenz-floss#373)

This allows most of the formula tests to be reliable.

Only the test for the `EVChargerPool.component_data` method still
depends on a lower level mock, and it needs it because it combines
resampled data with component state values. If it is flaky, we'll have
to add multi stage testing - with a fast test first, and a slower one if
that fails, etc.

Author: daniel-zullo-frequenz

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

@@ -110,3 +110,8 @@ def from_generator(
         ).generate()
         self._engines[channel_key] = engine
         return engine
+
+    async def stop(self) -> None:
+        """Stop all formula engines in the pool."""
+        for engine in self._engines.values():
+            await engine._stop()  # pylint: disable=protected-access

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

@@ -3,12 +3,20 @@
 
 """Formula generator for PV Power, from the component graph."""
 
+from __future__ import annotations
+
 import logging
+from collections import abc
 
 from ....microgrid import connection_manager
 from ....microgrid.component import ComponentCategory, ComponentMetricId, InverterType
 from .._formula_engine import FormulaEngine
-from ._formula_generator import NON_EXISTING_COMPONENT_ID, FormulaGenerator, FormulaType
+from ._formula_generator import (
+    NON_EXISTING_COMPONENT_ID,
+    FormulaGenerationError,
+    FormulaGenerator,
+    FormulaType,
+)
 
 _logger = logging.getLogger(__name__)
 
@@ -27,15 +35,8 @@ def generate(self) -> FormulaEngine:
         """
         builder = self._get_builder("pv-power", ComponentMetricId.ACTIVE_POWER)
 
-        component_graph = connection_manager.get().component_graph
-        pv_inverters = list(
-            comp
-            for comp in component_graph.components()
-            if comp.category == ComponentCategory.INVERTER
-            and comp.type == InverterType.SOLAR
-        )
-
-        if not pv_inverters:
+        pv_meters = self._get_pv_meters()
+        if not pv_meters:
             _logger.warning(
                 "Unable to find any PV inverters in the component graph. "
                 "Subscribing to the resampling actor with a non-existing "
@@ -51,11 +52,11 @@ def generate(self) -> FormulaEngine:
 
         builder.push_oper("(")
         builder.push_oper("(")
-        for idx, comp in enumerate(pv_inverters):
+        for idx, comp_id in enumerate(pv_meters):
             if idx > 0:
                 builder.push_oper("+")
 
-            builder.push_component_metric(comp.component_id, nones_are_zeros=True)
+            builder.push_component_metric(comp_id, nones_are_zeros=True)
         builder.push_oper(")")
         if self._config.formula_type == FormulaType.PRODUCTION:
             builder.push_oper("*")
@@ -66,3 +67,40 @@ def generate(self) -> FormulaEngine:
             builder.push_clipper(0.0, None)
 
         return builder.build()
+
+    def _get_pv_meters(self) -> abc.Set[int]:
+        component_graph = connection_manager.get().component_graph
+
+        pv_inverters = list(
+            comp
+            for comp in component_graph.components()
+            if comp.category == ComponentCategory.INVERTER
+            and comp.type == InverterType.SOLAR
+        )
+        pv_meters: set[int] = set()
+
+        if not pv_inverters:
+            return pv_meters
+
+        for pv_inverter in pv_inverters:
+            predecessors = component_graph.predecessors(pv_inverter.component_id)
+            if len(predecessors) != 1:
+                raise FormulaGenerationError(
+                    "Expected exactly one predecessor for PV inverter "
+                    f"{pv_inverter.component_id}, but found {len(predecessors)}."
+                )
+            meter = next(iter(predecessors))
+            if meter.category != ComponentCategory.METER:
+                raise FormulaGenerationError(
+                    f"Expected predecessor of PV inverter {pv_inverter.component_id} "
+                    f"to be a meter, but found {meter.category}."
+                )
+            meter_successors = component_graph.successors(meter.component_id)
+            if len(meter_successors) != 1:
+                raise FormulaGenerationError(
+                    f"Expected exactly one successor for meter {meter.component_id}"
+                    f", connected to PV inverter {pv_inverter.component_id}"
+                    f", but found {len(meter_successors)}."
+                )
+            pv_meters.add(meter.component_id)
+        return pv_meters

src/frequenz/sdk/timeseries/battery_pool/battery_pool.py

@@ -9,7 +9,7 @@
 import uuid
 from collections.abc import Set
 from datetime import timedelta
-from typing import Any
+from typing import Any, Awaitable
 
 from frequenz.channels import Receiver, Sender
 
@@ -78,6 +78,7 @@ def __init__(  # pylint: disable=too-many-arguments
 
         self._working_batteries: set[int] = set()
 
+        self._update_battery_status_task: asyncio.Task[None] | None = None
         if self._batteries:
             self._update_battery_status_task = asyncio.create_task(
                 self._update_battery_status(batteries_status_receiver)
@@ -266,10 +267,13 @@ async def power_bounds(
 
     async def stop(self) -> None:
         """Stop all pending async tasks."""
-        await asyncio.gather(
-            *[method.stop() for method in self._active_methods.values()],
-            cancel_and_await(self._update_battery_status_task),
-        )
+        tasks_to_stop: list[Awaitable[Any]] = [
+            method.stop() for method in self._active_methods.values()
+        ]
+        tasks_to_stop.append(self._formula_pool.stop())
+        if self._update_battery_status_task:
+            tasks_to_stop.append(cancel_and_await(self._update_battery_status_task))
+        await asyncio.gather(*tasks_to_stop)
 
     def _get_all_batteries(self) -> Set[int]:
         """Get all batteries from the microgrid.

src/frequenz/sdk/timeseries/ev_charger_pool/_ev_charger_pool.py

@@ -15,6 +15,7 @@
 
 from frequenz.channels import Broadcast, ChannelClosedError, Receiver, Sender
 
+from ..._internal._asyncio import cancel_and_await
 from ...actor import ChannelRegistry, ComponentMetricRequest
 from ...microgrid import connection_manager
 from ...microgrid.component import ComponentCategory, ComponentMetricId
@@ -271,6 +272,11 @@ async def stop(self) -> None:
             await self._bounds_setter.stop()
         if self._state_tracker:
             await self._state_tracker.stop()
+        await self._formula_pool.stop()
+        for stream in self._status_streams.values():
+            task, chan = stream
+            await chan.close()
+            await cancel_and_await(task)
 
     async def _get_current_streams(
         self, component_id: int

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

@@ -402,3 +402,7 @@ def chp_consumption_power(self) -> FormulaEngine:
         )
         assert isinstance(engine, FormulaEngine)
         return engine
+
+    async def stop(self) -> None:
+        """Stop all formula engines."""
+        await self._formula_pool.stop()

tests/timeseries/_battery_pool/test_battery_pool.py

@@ -13,7 +13,7 @@
 
 import async_solipsism
 import pytest
-from frequenz.channels import Broadcast, Receiver, Sender
+from frequenz.channels import Receiver, Sender
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
@@ -23,8 +23,7 @@
 )
 from frequenz.sdk.actor import ResamplerConfig
 from frequenz.sdk.actor.power_distributing import BatteryStatus
-from frequenz.sdk.microgrid.component import ComponentCategory, ComponentMetricId
-from frequenz.sdk.timeseries import Sample
+from frequenz.sdk.microgrid.component import ComponentCategory
 from frequenz.sdk.timeseries.battery_pool import (
     BatteryPool,
     Bound,
@@ -445,53 +444,24 @@ async def test_battery_pool_power(mocker: MockerFixture) -> None:
     """Test `BatteryPool.{,production,consumption}_power` methods."""
     mockgrid = MockMicrogrid(grid_side_meter=True)
     mockgrid.add_batteries(2)
-    await mockgrid.start(mocker)
-
-    channels: dict[int, Broadcast[Sample]] = {
-        meter_id: Broadcast(f"#{meter_id}")
-        for meter_id in [*mockgrid.meter_ids, *mockgrid.battery_inverter_ids]
-    }
-    senders: list[Sender[Sample]] = [
-        channels[component_id].new_sender()
-        for component_id in mockgrid.battery_inverter_ids
-    ]
-
-    async def send_resampled_data(
-        now: datetime,
-        meter_data: list[float | None],
-    ) -> None:
-        """Send resampled data to the channels."""
-        for sender, value in zip(senders, meter_data):
-            await sender.send(Sample(now, value))
-
-    def mock_resampled_receiver(
-        _1: Any, component_id: int, _2: ComponentMetricId
-    ) -> Receiver[Sample]:
-        return channels[component_id].new_receiver()
-
-    mocker.patch(
-        "frequenz.sdk.timeseries._formula_engine._resampled_formula_builder"
-        ".ResampledFormulaBuilder._get_resampled_receiver",
-        mock_resampled_receiver,
-    )
+    await mockgrid.start_mock_datapipeline(mocker)
 
     battery_pool = microgrid.battery_pool()
     power_receiver = battery_pool.power.new_receiver()
     consumption_receiver = battery_pool.consumption_power.new_receiver()
     production_receiver = battery_pool.production_power.new_receiver()
 
-    now = datetime.now(tz=timezone.utc)
-    await send_resampled_data(now, [2.0, 3.0])
+    await mockgrid.mock_data.send_bat_inverter_power([2.0, 3.0])
     assert (await power_receiver.receive()).value == 5.0
     assert (await consumption_receiver.receive()).value == 5.0
     assert (await production_receiver.receive()).value == 0.0
 
-    await send_resampled_data(now + timedelta(seconds=1), [-2.0, -5.0])
+    await mockgrid.mock_data.send_bat_inverter_power([-2.0, -5.0])
     assert (await power_receiver.receive()).value == -7.0
     assert (await consumption_receiver.receive()).value == 0.0
     assert (await production_receiver.receive()).value == 7.0
 
-    await send_resampled_data(now + timedelta(seconds=3), [2.0, -5.0])
+    await mockgrid.mock_data.send_bat_inverter_power([2.0, -5.0])
     assert (await power_receiver.receive()).value == -3.0
     assert (await consumption_receiver.receive()).value == 0.0
     assert (await production_receiver.receive()).value == 3.0