Use the MockMicrogrid as an async context manager

This commit updates many tests that use the `MockMicrogrid` and finalize
it manually to use it as a context manager so finalization is done
automatically.

This commit can be reviewed much more easily hiding differences in
spaces (with `diff -w`, or appending `?w=1` to the GitHub URL).

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

tests/actor/power_distributing/test_power_distributing.py

@@ -80,10 +80,10 @@ async def new(
         grid_meter: bool | None = None,
     ) -> _Mocks:
         """Initialize the mocks."""
-        mockgrid = MockMicrogrid(graph=graph, grid_meter=grid_meter)
+        mockgrid = MockMicrogrid(graph=graph, grid_meter=grid_meter, mocker=mocker)
         if not graph:
             mockgrid.add_batteries(3)
-        await mockgrid.start(mocker)
+        await mockgrid.start()
 
         # pylint: disable=protected-access
         if microgrid._data_pipeline._DATA_PIPELINE is not None:
@@ -154,56 +154,59 @@ async def _patch_battery_pool_status(
             )
         )
 
-    async def test_constructor(self, mocker: MockerFixture) -> None:
-        """Test if gets all necessary data."""
-        mockgrid = MockMicrogrid(grid_meter=True)
+    async def test_constructor_with_grid_meter(self, mocker: MockerFixture) -> None:
+        """Test the constructor works with a grid meter."""
+        mockgrid = MockMicrogrid(grid_meter=True, mocker=mocker)
         mockgrid.add_batteries(2)
         mockgrid.add_batteries(1, no_meter=True)
-        await mockgrid.start(mocker)
 
-        requests_channel = Broadcast[Request]("power_distributor requests")
-        results_channel = Broadcast[Result]("power_distributor results")
-        battery_status_channel = Broadcast[ComponentPoolStatus]("battery_status")
-        async with PowerDistributingActor(
-            requests_receiver=requests_channel.new_receiver(),
-            results_sender=results_channel.new_sender(),
-            component_pool_status_sender=battery_status_channel.new_sender(),
-        ) as distributor:
-            assert isinstance(distributor._component_manager, BatteryManager)
-            assert distributor._component_manager._bat_invs_map == {
-                9: {8},
-                19: {18},
-                29: {28},
-            }
-            assert distributor._component_manager._inv_bats_map == {
-                8: {9},
-                18: {19},
-                28: {29},
-            }
-        await mockgrid.cleanup()
-
-        # Test if it works without grid side meter
-        mockgrid = MockMicrogrid(grid_meter=False)
+        async with mockgrid:
+            requests_channel = Broadcast[Request]("power_distributor requests")
+            results_channel = Broadcast[Result]("power_distributor results")
+            battery_status_channel = Broadcast[ComponentPoolStatus]("battery_status")
+            async with PowerDistributingActor(
+                requests_receiver=requests_channel.new_receiver(),
+                results_sender=results_channel.new_sender(),
+                component_pool_status_sender=battery_status_channel.new_sender(),
+            ) as distributor:
+                assert isinstance(distributor._component_manager, BatteryManager)
+                assert distributor._component_manager._bat_invs_map == {
+                    9: {8},
+                    19: {18},
+                    29: {28},
+                }
+                assert distributor._component_manager._inv_bats_map == {
+                    8: {9},
+                    18: {19},
+                    28: {29},
+                }
+
+    async def test_constructor_without_grid_meter(self, mocker: MockerFixture) -> None:
+        """Test the constructor works without a grid meter."""
+        mockgrid = MockMicrogrid(grid_meter=False, mocker=mocker)
         mockgrid.add_batteries(1)
         mockgrid.add_batteries(2, no_meter=True)
-        await mockgrid.start(mocker)
-        async with PowerDistributingActor(
-            requests_receiver=requests_channel.new_receiver(),
-            results_sender=results_channel.new_sender(),
-            component_pool_status_sender=battery_status_channel.new_sender(),
-        ) as distributor:
-            assert isinstance(distributor._component_manager, BatteryManager)
-            assert distributor._component_manager._bat_invs_map == {
-                9: {8},
-                19: {18},
-                29: {28},
-            }
-            assert distributor._component_manager._inv_bats_map == {
-                8: {9},
-                18: {19},
-                28: {29},
-            }
-        await mockgrid.cleanup()
+
+        async with mockgrid:
+            requests_channel = Broadcast[Request]("power_distributor requests")
+            results_channel = Broadcast[Result]("power_distributor results")
+            battery_status_channel = Broadcast[ComponentPoolStatus]("battery_status")
+            async with PowerDistributingActor(
+                requests_receiver=requests_channel.new_receiver(),
+                results_sender=results_channel.new_sender(),
+                component_pool_status_sender=battery_status_channel.new_sender(),
+            ) as distributor:
+                assert isinstance(distributor._component_manager, BatteryManager)
+                assert distributor._component_manager._bat_invs_map == {
+                    9: {8},
+                    19: {18},
+                    29: {28},
+                }
+                assert distributor._component_manager._inv_bats_map == {
+                    8: {9},
+                    18: {19},
+                    28: {29},
+                }
 
     async def init_component_data(
         self,

tests/actor/test_battery_pool_status.py

@@ -32,77 +32,85 @@ async def test_batteries_status(self, mocker: MockerFixture) -> None:
         Args:
             mocker: Pytest mocker fixture.
         """
-        mock_microgrid = MockMicrogrid(grid_meter=True)
+        mock_microgrid = MockMicrogrid(grid_meter=True, mocker=mocker)
         mock_microgrid.add_batteries(3)
-        await mock_microgrid.start(mocker)
 
-        batteries = {
-            battery.component_id
-            for battery in mock_microgrid.mock_client.component_graph.components(
-                component_categories={ComponentCategory.BATTERY}
+        async with mock_microgrid:
+            batteries = {
+                battery.component_id
+                for battery in mock_microgrid.mock_client.component_graph.components(
+                    component_categories={ComponentCategory.BATTERY}
+                )
+            }
+            battery_status_channel = Broadcast[ComponentPoolStatus]("battery_status")
+            battery_status_recv = battery_status_channel.new_receiver(maxsize=1)
+            batteries_status = ComponentPoolStatusTracker(
+                component_ids=batteries,
+                component_status_sender=battery_status_channel.new_sender(),
+                max_data_age_sec=5,
+                max_blocking_duration_sec=30,
+                component_status_tracker_type=BatteryStatusTracker,
             )
-        }
-        battery_status_channel = Broadcast[ComponentPoolStatus]("battery_status")
-        battery_status_recv = battery_status_channel.new_receiver(maxsize=1)
-        batteries_status = ComponentPoolStatusTracker(
-            component_ids=batteries,
-            component_status_sender=battery_status_channel.new_sender(),
-            max_data_age_sec=5,
-            max_blocking_duration_sec=30,
-            component_status_tracker_type=BatteryStatusTracker,
-        )
-        await asyncio.sleep(0.1)
-
-        expected_working: set[int] = set()
-        assert batteries_status.get_working_components(batteries) == expected_working
-
-        batteries_list = list(batteries)
-
-        await mock_microgrid.mock_client.send(
-            battery_data(component_id=batteries_list[0])
-        )
-        await asyncio.sleep(0.1)
-        assert batteries_status.get_working_components(batteries) == expected_working
-
-        expected_working.add(batteries_list[0])
-        await mock_microgrid.mock_client.send(
-            inverter_data(component_id=batteries_list[0] - 1)
-        )
-        await asyncio.sleep(0.1)
-        assert batteries_status.get_working_components(batteries) == expected_working
-        msg = await asyncio.wait_for(battery_status_recv.receive(), timeout=0.2)
-        assert msg == batteries_status._current_status
-
-        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.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)
-        assert batteries_status.get_working_components(batteries) == expected_working
-        msg = await asyncio.wait_for(battery_status_recv.receive(), timeout=0.2)
-        assert msg == batteries_status._current_status
-
-        await batteries_status.update_status(
-            succeeded_components={9}, failed_components={19, 29}
-        )
-        await asyncio.sleep(0.1)
-        assert batteries_status.get_working_components(batteries) == {9}
-
-        await batteries_status.update_status(
-            succeeded_components={9, 19}, failed_components=set()
-        )
-        await asyncio.sleep(0.1)
-        assert batteries_status.get_working_components(batteries) == {9, 19}
-
-        await batteries_status.stop()
+            await asyncio.sleep(0.1)
+
+            expected_working: set[int] = set()
+            assert (
+                batteries_status.get_working_components(batteries) == expected_working
+            )
+
+            batteries_list = list(batteries)
+
+            await mock_microgrid.mock_client.send(
+                battery_data(component_id=batteries_list[0])
+            )
+            await asyncio.sleep(0.1)
+            assert (
+                batteries_status.get_working_components(batteries) == expected_working
+            )
+
+            expected_working.add(batteries_list[0])
+            await mock_microgrid.mock_client.send(
+                inverter_data(component_id=batteries_list[0] - 1)
+            )
+            await asyncio.sleep(0.1)
+            assert (
+                batteries_status.get_working_components(batteries) == expected_working
+            )
+            msg = await asyncio.wait_for(battery_status_recv.receive(), timeout=0.2)
+            assert msg == batteries_status._current_status
+
+            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.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)
+            assert (
+                batteries_status.get_working_components(batteries) == expected_working
+            )
+            msg = await asyncio.wait_for(battery_status_recv.receive(), timeout=0.2)
+            assert msg == batteries_status._current_status
+
+            await batteries_status.update_status(
+                succeeded_components={9}, failed_components={19, 29}
+            )
+            await asyncio.sleep(0.1)
+            assert batteries_status.get_working_components(batteries) == {9}
+
+            await batteries_status.update_status(
+                succeeded_components={9, 19}, failed_components=set()
+            )
+            await asyncio.sleep(0.1)
+            assert batteries_status.get_working_components(batteries) == {9, 19}
+
+            await batteries_status.stop()