WIP: Fix tests

Author: llucax

tests/timeseries/_battery_pool/test_battery_pool_control_methods.py

@@ -5,13 +5,15 @@
 
 import asyncio
 import dataclasses
-import typing
+import re
+from collections.abc import AsyncIterator, Callable
 from datetime import datetime, timedelta, timezone
+from typing import cast
 from unittest.mock import AsyncMock, MagicMock
 
 import async_solipsism
 import pytest
-from frequenz.channels import LatestValueCache, Sender
+from frequenz.channels import LatestValueCache, Receiver, Sender
 from frequenz.quantities import Power
 from pytest_mock import MockerFixture
 
@@ -54,7 +56,7 @@ class Mocks:
 
 
 @pytest.fixture
-async def mocks(mocker: MockerFixture) -> typing.AsyncIterator[Mocks]:
+async def mocks(mocker: MockerFixture) -> AsyncIterator[Mocks]:
     """Fixture for the mocks."""
     mockgrid = MockMicrogrid()
     mockgrid.add_batteries(4)
@@ -71,19 +73,15 @@ async def mocks(mocker: MockerFixture) -> typing.AsyncIterator[Mocks]:
     dp = microgrid._data_pipeline._DATA_PIPELINE
     assert dp is not None
 
-    yield Mocks(
+    _mocks = Mocks(
         mockgrid,
         streamer,
         dp._battery_power_wrapper.status_channel.new_sender(),
     )
-
-    await asyncio.gather(
-        *[
-            dp._stop(),
-            streamer.stop(),
-            mockgrid.cleanup(),
-        ]
-    )
+    try:
+        yield _mocks
+    finally:
+        await asyncio.gather(dp._stop(), streamer.stop(), mockgrid.cleanup())
 
 
 class TestBatteryPoolControl:
@@ -159,16 +157,17 @@ async def _init_data_for_inverters(self, mocks: Mocks) -> None:
 
     def _assert_report(  # pylint: disable=too-many-arguments
         self,
-        report: BatteryPoolReport,
+        report: BatteryPoolReport | None,
         *,
         power: float | None,
         lower: float,
         upper: float,
         dist_result: _power_distributing.Result | None = None,
         expected_result_pred: (
-            typing.Callable[[_power_distributing.Result], bool] | None
+            Callable[[_power_distributing.Result], bool] | None
         ) = None,
     ) -> None:
+        assert report is not None
         assert report.target_power == (
             Power.from_watts(power) if power is not None else None
         )
@@ -179,6 +178,22 @@ def _assert_report(  # pylint: disable=too-many-arguments
             assert dist_result is not None
             assert expected_result_pred(dist_result)
 
+    async def _recv_reports_until(
+        self,
+        bounds_rx: Receiver[BatteryPoolReport],
+        check: Callable[[BatteryPoolReport], bool],
+    ) -> BatteryPoolReport | None:
+        """Receive reports until the given condition is met."""
+        max_reports = 10
+        ctr = 0
+        while ctr < max_reports:
+            ctr += 1
+            async with asyncio.timeout(10.0):
+                report = await bounds_rx.receive()
+            if check(report):
+                return report
+        return None
+
     async def test_case_1(
         self,
         mocks: Mocks,
@@ -189,7 +204,7 @@ async def test_case_1(
         - single battery pool with all batteries.
         - all batteries are working, then one battery stops working.
         """
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )
@@ -205,9 +220,12 @@ async def test_case_1(
             battery_pool.power_distribution_results.new_receiver()
         )
 
-        self._assert_report(
-            await bounds_rx.receive(), power=None, lower=-4000.0, upper=4000.0
+        report = await self._recv_reports_until(
+            bounds_rx,
+            lambda r: r.bounds is not None
+            and r.bounds.upper == Power.from_watts(4000.0),
         )
+        self._assert_report(report, power=None, lower=-4000.0, upper=4000.0)
 
         await battery_pool.propose_power(Power.from_watts(1000.0))
 
@@ -293,7 +311,7 @@ async def test_case_2(self, mocks: Mocks, mocker: MockerFixture) -> None:
         - two battery pools with different batteries.
         - all batteries are working.
         """
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )
@@ -351,7 +369,7 @@ async def test_case_3(self, mocks: Mocks, mocker: MockerFixture) -> None:
         - two battery pools with same batteries, but different priorities.
         - all batteries are working.
         """
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )
@@ -415,7 +433,7 @@ async def test_case_4(self, mocks: Mocks, mocker: MockerFixture) -> None:
         - single battery pool with all batteries.
         - all batteries are working, but have exclusion bounds.
         """
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )
@@ -444,6 +462,7 @@ async def test_case_4(self, mocks: Mocks, mocker: MockerFixture) -> None:
             mocker.call(inv_id, 250.0)
             for inv_id in mocks.microgrid.battery_inverter_ids
         ]
+
         self._assert_report(
             await bounds_rx.receive(),
             power=1000.0,
@@ -461,9 +480,10 @@ async def test_case_4(self, mocks: Mocks, mocker: MockerFixture) -> None:
         # available power.
         await battery_pool.propose_power(Power.from_watts(50.0))
 
-        self._assert_report(
-            await bounds_rx.receive(), power=400.0, lower=-4000.0, upper=4000.0
+        report = await self._recv_reports_until(
+            bounds_rx, lambda r: r.target_power == Power.from_watts(400.0)
         )
+        self._assert_report(report, power=400.0, lower=-4000.0, upper=4000.0)
         await asyncio.sleep(0.0)  # Wait for the power to be distributed.
         assert set_power.call_count == 4
         assert sorted(set_power.call_args_list) == [
@@ -542,7 +562,7 @@ async def test_case_5(  # pylint: disable=too-many-statements,too-many-locals
         - two battery pools are in the shifting group, two are not.
         - all batteries are working.
         """
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )

tests/timeseries/_ev_charger_pool/test_ev_charger_pool_control_methods.py

@@ -4,8 +4,9 @@
 """Test the EV charger pool control methods."""
 
 import asyncio
-import typing
+from collections.abc import AsyncIterator, Callable
 from datetime import datetime, timedelta, timezone
+from typing import cast
 from unittest.mock import AsyncMock, MagicMock
 
 import async_solipsism
@@ -15,6 +16,7 @@
 from frequenz.client.microgrid.component import ComponentStateCode
 from frequenz.quantities import Power, Voltage
 from pytest_mock import MockerFixture
+from regex import E
 
 from frequenz.sdk import microgrid
 from frequenz.sdk.microgrid import _power_distributing
@@ -41,7 +43,7 @@ def event_loop_policy() -> async_solipsism.EventLoopPolicy:
 
 
 @pytest.fixture
-async def mocks(mocker: MockerFixture) -> typing.AsyncIterator[_Mocks]:
+async def mocks(mocker: MockerFixture) -> AsyncIterator[_Mocks]:
     """Create the mocks."""
     mockgrid = MockMicrogrid(grid_meter=True)
     mockgrid.add_ev_chargers(4)
@@ -55,13 +57,17 @@ async def mocks(mocker: MockerFixture) -> typing.AsyncIterator[_Mocks]:
     )
     streamer = MockComponentDataStreamer(mockgrid.mock_client)
 
-    dp = typing.cast(_DataPipeline, microgrid._data_pipeline._DATA_PIPELINE)
+    dp = cast(_DataPipeline, microgrid._data_pipeline._DATA_PIPELINE)
 
-    yield _Mocks(
+    _mocks = _Mocks(
         mockgrid,
         streamer,
         dp._ev_power_wrapper.status_channel.new_sender(),
     )
+    try:
+        yield _mocks
+    finally:
+        await _mocks.stop()
 
 
 class TestEVChargerPoolControl:
@@ -145,16 +151,17 @@ async def _init_ev_chargers(self, mocks: _Mocks) -> None:
 
     def _assert_report(  # pylint: disable=too-many-arguments
         self,
-        report: EVChargerPoolReport,
+        report: EVChargerPoolReport | None,
         *,
         power: float | None,
         lower: float,
         upper: float,
         dist_result: _power_distributing.Result | None = None,
         expected_result_pred: (
-            typing.Callable[[_power_distributing.Result], bool] | None
+            Callable[[_power_distributing.Result], bool] | None
         ) = None,
     ) -> None:
+        assert report is not None
         assert report.target_power == (
             Power.from_watts(power) if power is not None else None
         )
@@ -165,23 +172,21 @@ def _assert_report(  # pylint: disable=too-many-arguments
             assert dist_result is not None
             assert expected_result_pred(dist_result)
 
-    async def _get_bounds_receiver(
-        self, ev_charger_pool: EVChargerPool
-    ) -> Receiver[EVChargerPoolReport]:
-        bounds_rx = ev_charger_pool.power_status.new_receiver()
-
-        # Consume initial reports as chargers are initialized
-        expected_upper_bounds = 44160.0
+    async def _recv_reports_until(
+        self,
+        bounds_rx: Receiver[EVChargerPoolReport],
+        check: Callable[[EVChargerPoolReport], bool],
+    ) -> EVChargerPoolReport | None:
+        """Receive reports until the given condition is met."""
         max_reports = 10
         ctr = 0
         while ctr < max_reports:
             ctr += 1
-            report = await bounds_rx.receive()
-            assert report.bounds is not None
-            if report.bounds.upper == Power.from_watts(expected_upper_bounds):
-                break
-
-        return bounds_rx
+            async with asyncio.timeout(10.0):
+                report = await bounds_rx.receive()
+            if check(report):
+                return report
+        return None
 
     async def test_setting_power(
         self,
@@ -192,7 +197,7 @@ async def test_setting_power(
         traveller = time_machine.travel(datetime(2012, 12, 12))
         mock_time = traveller.start()
 
-        set_power = typing.cast(
+        set_power = cast(
             AsyncMock,
             microgrid.connection_manager.get().api_client.set_component_power_active,
         )
@@ -201,23 +206,27 @@ async def test_setting_power(
         await self._patch_ev_pool_status(mocks, mocker)
         await self._patch_power_distributing_actor(mocker)
 
-        bounds_rx = await self._get_bounds_receiver(ev_charger_pool)
+        bounds_rx = ev_charger_pool.power_status.new_receiver()
+        # Receive reports until all chargers are initialized
+        report = await self._recv_reports_until(
+            bounds_rx,
+            lambda r: r.bounds is not None
+            and r.bounds.upper == Power.from_watts(44160.0),
+        )
 
         # Check that chargers are initialized to Power.zero()
         assert set_power.call_count == 4
         assert all(x.args[1] == 0.0 for x in set_power.call_args_list)
-
-        self._assert_report(
-            await bounds_rx.receive(), power=None, lower=0.0, upper=44160.0
-        )
+        self._assert_report(report, power=None, lower=0.0, upper=44160.0)
 
         set_power.reset_mock()
         await ev_charger_pool.propose_power(Power.from_watts(40000.0))
         # ignore one report because it is not always immediately updated.
-        await bounds_rx.receive()
-        self._assert_report(
-            await bounds_rx.receive(), power=40000.0, lower=0.0, upper=44160.0
+        report = await self._recv_reports_until(
+            bounds_rx,
+            lambda r: r.target_power == Power.from_watts(40000.0),
         )
+        self._assert_report(report, power=40000.0, lower=0.0, upper=44160.0)
         mock_time.shift(timedelta(seconds=60))
         await asyncio.sleep(0.15)
 
@@ -240,7 +249,7 @@ async def test_setting_power(
         # Throttle the power
         set_power.reset_mock()
         await ev_charger_pool.propose_power(Power.from_watts(32000.0))
-        await bounds_rx.receive()
+        report = await bounds_rx.receive()
         await asyncio.sleep(0.02)
         assert set_power.call_count == 1