Update PowerDistributingActor tests to use MockMicrogrid

MockMicrogrid makes it easy to dynamically create component graph
configurations to test with, rather than having to use manually
written component graphs with the lower level MockMicrogridClient.

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

Author: shsms

tests/actor/test_battery_pool_status.py

@@ -14,6 +14,7 @@
     BatteryStatus,
 )
 from frequenz.sdk.microgrid.component import ComponentCategory
+from tests.timeseries.mock_microgrid import MockMicrogrid
 
 from ..utils.mock_microgrid_client import MockMicrogridClient
 from .test_battery_status import battery_data, component_graph, inverter_data
@@ -38,17 +39,21 @@ async def mock_microgrid(self, mocker: MockerFixture) -> MockMicrogridClient:
         microgrid.initialize(mocker)
         return microgrid
 
-    async def test_batteries_status(self, mock_microgrid: MockMicrogridClient) -> None:
+    async def test_batteries_status(self, mocker: MockerFixture) -> None:
         """Basic tests for BatteryPoolStatus.
 
         BatteryStatusTracker is more tested in its own unit tests.
 
         Args:
             mock_microgrid: mock microgrid client
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
+
         batteries = {
             battery.component_id
-            for battery in mock_microgrid.component_graph.components(
+            for battery in mock_microgrid.mock_client.component_graph.components(
                 component_category={ComponentCategory.BATTERY}
             )
         }
@@ -67,22 +72,34 @@ async def test_batteries_status(self, mock_microgrid: MockMicrogridClient) -> No
 
         batteries_list = list(batteries)
 
-        await mock_microgrid.send(battery_data(component_id=batteries_list[0]))
+        await mock_microgrid.mock_client.send(
+            battery_data(component_id=batteries_list[0])
+        )
         await asyncio.sleep(0.1)
         assert batteries_status.get_working_batteries(batteries) == expected_working
 
         expected_working.add(batteries_list[0])
-        await mock_microgrid.send(inverter_data(component_id=batteries_list[0] - 1))
+        await mock_microgrid.mock_client.send(
+            inverter_data(component_id=batteries_list[0] - 1)
+        )
         await asyncio.sleep(0.1)
         assert batteries_status.get_working_batteries(batteries) == expected_working
         msg = await asyncio.wait_for(battery_status_recv.receive(), timeout=0.2)
         assert msg == batteries_status._current_status
 
-        await mock_microgrid.send(inverter_data(component_id=batteries_list[1] - 1))
-        await mock_microgrid.send(battery_data(component_id=batteries_list[1]))
+        await mock_microgrid.mock_client.send(
+            inverter_data(component_id=batteries_list[1] - 1)
+        )
+        await mock_microgrid.mock_client.send(
+            battery_data(component_id=batteries_list[1])
+        )
 
-        await mock_microgrid.send(inverter_data(component_id=batteries_list[2] - 1))
-        await mock_microgrid.send(battery_data(component_id=batteries_list[2]))
+        await mock_microgrid.mock_client.send(
+            inverter_data(component_id=batteries_list[2] - 1)
+        )
+        await mock_microgrid.mock_client.send(
+            battery_data(component_id=batteries_list[2])
+        )
 
         expected_working = set(batteries_list)
         await asyncio.sleep(0.1)
@@ -91,15 +108,15 @@ async def test_batteries_status(self, mock_microgrid: MockMicrogridClient) -> No
         assert msg == batteries_status._current_status
 
         await batteries_status.update_status(
-            succeed_batteries={106}, failed_batteries={206, 306}
+            succeed_batteries={9}, failed_batteries={19, 29}
         )
         await asyncio.sleep(0.1)
-        assert batteries_status.get_working_batteries(batteries) == {106}
+        assert batteries_status.get_working_batteries(batteries) == {9}
 
         await batteries_status.update_status(
-            succeed_batteries={106, 206}, failed_batteries=set()
+            succeed_batteries={9, 19}, failed_batteries=set()
         )
         await asyncio.sleep(0.1)
-        assert batteries_status.get_working_batteries(batteries) == {106, 206}
+        assert batteries_status.get_working_batteries(batteries) == {9, 19}
 
         await batteries_status.stop()

tests/actor/test_battery_status.py

@@ -33,6 +33,7 @@
     ComponentCategory,
     InverterData,
 )
+from tests.timeseries.mock_microgrid import MockMicrogrid
 
 from ..utils.component_data_wrapper import BatteryDataWrapper, InverterDataWrapper
 from ..utils.mock_microgrid_client import MockMicrogridClient
@@ -154,8 +155,8 @@ class Message(Generic[T]):
     inner: T
 
 
-BATTERY_ID = 106
-INVERTER_ID = 105
+BATTERY_ID = 9
+INVERTER_ID = 8
 
 
 # pylint: disable=protected-access, unused-argument
@@ -179,7 +180,7 @@ async def mock_microgrid(self, mocker: MockerFixture) -> MockMicrogridClient:
 
     @time_machine.travel("2022-01-01 00:00 UTC", tick=False)
     async def test_sync_update_status_with_messages(
-        self, mock_microgrid: MockMicrogridClient
+        self, mocker: MockerFixture
     ) -> None:
         """Test if messages changes battery status/
 
@@ -189,6 +190,10 @@ async def test_sync_update_status_with_messages(
         Args:
             mock_microgrid: mock_microgrid fixture
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
+
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
@@ -336,10 +341,9 @@ async def test_sync_update_status_with_messages(
         assert tracker._get_new_status_if_changed() is Status.NOT_WORKING
 
         await tracker.stop()
+        await mock_microgrid.cleanup()
 
-    async def test_sync_blocking_feature(
-        self, mock_microgrid: MockMicrogridClient
-    ) -> None:
+    async def test_sync_blocking_feature(self, mocker: MockerFixture) -> None:
         """Test if status changes when SetPowerResult message is received.
 
         Tests uses FakeSelect to test status in sync way.
@@ -348,6 +352,9 @@ async def test_sync_blocking_feature(
         Args:
             mock_microgrid: mock_microgrid fixture
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
 
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
@@ -444,14 +451,15 @@ async def test_sync_blocking_feature(
 
             # If battery succeed, then it should unblock.
             tracker._handle_status_set_power_result(
-                SetPowerResult(succeed={BATTERY_ID}, failed={206})
+                SetPowerResult(succeed={BATTERY_ID}, failed={19})
             )
             assert tracker._get_new_status_if_changed() is Status.WORKING
 
         await tracker.stop()
+        await mock_microgrid.cleanup()
 
     async def test_sync_blocking_interrupted_with_with_max_data(
-        self, mock_microgrid: MockMicrogridClient
+        self, mocker: MockerFixture
     ) -> None:
         """Test if status changes when SetPowerResult message is received.
 
@@ -461,6 +469,9 @@ async def test_sync_blocking_interrupted_with_with_max_data(
         Args:
             mock_microgrid: mock_microgrid fixture
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
 
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
@@ -498,7 +509,7 @@ async def test_sync_blocking_interrupted_with_with_max_data(
 
     @time_machine.travel("2022-01-01 00:00 UTC", tick=False)
     async def test_sync_blocking_interrupted_with_invalid_message(
-        self, mock_microgrid: MockMicrogridClient
+        self, mocker: MockerFixture
     ) -> None:
         """Test if status changes when SetPowerResult message is received.
 
@@ -508,6 +519,9 @@ async def test_sync_blocking_interrupted_with_invalid_message(
         Args:
             mock_microgrid: mock_microgrid fixture
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
 
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
@@ -555,9 +569,7 @@ async def test_sync_blocking_interrupted_with_invalid_message(
         await tracker.stop()
 
     @time_machine.travel("2022-01-01 00:00 UTC", tick=False)
-    async def test_timers(
-        self, mock_microgrid: MockMicrogridClient, mocker: MockerFixture
-    ) -> None:
+    async def test_timers(self, mocker: MockerFixture) -> None:
         """Test if messages changes battery status/
 
         Tests uses FakeSelect to test status in sync way.
@@ -567,6 +579,10 @@ async def test_timers(
             mock_microgrid: mock_microgrid fixture
             mocker: pytest mocker instance
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
+
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
@@ -616,16 +632,18 @@ async def test_timers(
 
         assert inverter_timer_spy.call_count == 2
         await tracker.stop()
+        await mock_microgrid.cleanup()
 
     @time_machine.travel("2022-01-01 00:00 UTC", tick=False)
-    async def test_async_battery_status(
-        self, mock_microgrid: MockMicrogridClient
-    ) -> None:
+    async def test_async_battery_status(self, mocker: MockerFixture) -> None:
         """Test if status changes.
 
         Args:
             mock_microgrid: mock_microgrid fixture
         """
+        mock_microgrid = MockMicrogrid(grid_side_meter=True)
+        mock_microgrid.add_batteries(3)
+        await mock_microgrid.start(mocker)
 
         status_channel = Broadcast[Status]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
@@ -643,8 +661,10 @@ async def test_async_battery_status(
         await asyncio.sleep(0.01)
 
         with time_machine.travel("2022-01-01 00:00 UTC", tick=False) as time:
-            await mock_microgrid.send(inverter_data(component_id=INVERTER_ID))
-            await mock_microgrid.send(battery_data(component_id=BATTERY_ID))
+            await mock_microgrid.mock_client.send(
+                inverter_data(component_id=INVERTER_ID)
+            )
+            await mock_microgrid.mock_client.send(battery_data(component_id=BATTERY_ID))
             status = await asyncio.wait_for(status_receiver.receive(), timeout=0.1)
             assert status is Status.WORKING
 
@@ -656,11 +676,11 @@ async def test_async_battery_status(
 
             time.shift(2)
 
-            await mock_microgrid.send(battery_data(component_id=BATTERY_ID))
+            await mock_microgrid.mock_client.send(battery_data(component_id=BATTERY_ID))
             status = await asyncio.wait_for(status_receiver.receive(), timeout=0.1)
             assert status is Status.WORKING
 
-            await mock_microgrid.send(
+            await mock_microgrid.mock_client.send(
                 inverter_data(
                     component_id=INVERTER_ID,
                     timestamp=datetime.now(tz=timezone.utc) - timedelta(seconds=7),
@@ -675,8 +695,11 @@ async def test_async_battery_status(
             await asyncio.sleep(0.3)
             assert len(status_receiver) == 0
 
-            await mock_microgrid.send(inverter_data(component_id=INVERTER_ID))
+            await mock_microgrid.mock_client.send(
+                inverter_data(component_id=INVERTER_ID)
+            )
             status = await asyncio.wait_for(status_receiver.receive(), timeout=0.1)
             assert status is Status.WORKING
 
         await tracker.stop()
+        await mock_microgrid.cleanup()