Mock microgrid.get to return MockMicrogrid

That is optional, because in some cases it might not be needed.

Signed-off-by: ela-kotulska-frequenz <[email protected]>

Author: ela-kotulska-frequenz

tests/actor/test_power_distributing.py

@@ -6,7 +6,6 @@
 import re
 from dataclasses import dataclass
 from typing import Set, Tuple, TypeVar
-from unittest import mock
 from unittest.mock import AsyncMock, MagicMock
 
 from frequenz.channels import Bidirectional, Receiver, Sender
@@ -81,31 +80,30 @@ def component_graph(self) -> Tuple[Set[Component], Set[Connection]]:
         }
         return components, connections
 
-    async def test_constructor(self) -> None:
+    async def test_constructor(self, mocker: MockerFixture) -> None:
         """Test if gets all necessary data."""
         components, connections = self.component_graph()
         mock_microgrid = MockMicrogridClient(components, connections)
+        await mock_microgrid.initialize(mocker)
 
-        with mock.patch(
-            "frequenz.sdk.microgrid.get", return_value=mock_microgrid.mock_microgrid
-        ):
-            channel = Bidirectional[Request, Result]("user1", "power_distributor")
-            distributor = PowerDistributingActor({"user1": channel.service_handle})
+        channel = Bidirectional[Request, Result]("user1", "power_distributor")
+        distributor = PowerDistributingActor({"user1": channel.service_handle})
 
-            assert distributor._bat_inv_map == {106: 105, 206: 205, 306: 305}
-            assert distributor._inv_bat_map == {105: 106, 205: 206, 305: 306}
-            await distributor._stop_actor()
+        assert distributor._bat_inv_map == {106: 105, 206: 205, 306: 305}
+        assert distributor._inv_bat_map == {105: 106, 205: 206, 305: 306}
+        await distributor._stop_actor()
 
-    async def init_mock_microgrid(self) -> MockMicrogridClient:
+    async def init_mock_microgrid(self, mocker: MockerFixture) -> MockMicrogridClient:
         """Create mock microgrid and send initial data from the components.
 
         Returns:
             Mock microgrid instance.
         """
         components, connections = self.component_graph()
         microgrid = MockMicrogridClient(components, connections)
-        graph = microgrid.component_graph
+        await microgrid.initialize(mocker)
 
+        graph = microgrid.component_graph
         for battery in graph.components(component_category={ComponentCategory.BATTERY}):
             assert await microgrid.send(
                 battery_msg(
@@ -130,7 +128,7 @@ async def init_mock_microgrid(self) -> MockMicrogridClient:
     async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works with single user works."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel = Bidirectional[Request, Result]("user1", "power_distributor")
 
@@ -141,10 +139,6 @@ async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
-
         distributor = PowerDistributingActor({"user1": channel.service_handle})
 
         # Mock that all requested batteries are working.
@@ -172,7 +166,7 @@ async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
     async def test_power_distributor_two_users(self, mocker: MockerFixture) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works with two users."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         channel2 = Bidirectional[Request, Result]("user2", "power_distributor")
@@ -182,9 +176,6 @@ async def test_power_distributor_two_users(self, mocker: MockerFixture) -> None:
         }
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -228,7 +219,7 @@ async def test_power_distributor_invalid_battery_id(
     ) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution raises error if any battery id is invalid."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         service_channels = {
@@ -240,9 +231,6 @@ async def test_power_distributor_invalid_battery_id(
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -271,7 +259,7 @@ async def test_power_distributor_overlapping_batteries(
     ) -> None:
         # pylint: disable=too-many-locals
         """Test if requests with overlapping set of batteries are processed."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         channel2 = Bidirectional[Request, Result]("user2", "power_distributor")
@@ -283,9 +271,6 @@ async def test_power_distributor_overlapping_batteries(
         }
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -350,7 +335,7 @@ async def test_power_distributor_one_user_adjust_power_consume(
     ) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works with single user works."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         service_channels = {
@@ -365,9 +350,6 @@ async def test_power_distributor_one_user_adjust_power_consume(
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -398,7 +380,7 @@ async def test_power_distributor_one_user_adjust_power_supply(
     ) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works with single user works."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         service_channels = {
@@ -413,9 +395,6 @@ async def test_power_distributor_one_user_adjust_power_supply(
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -446,7 +425,7 @@ async def test_power_distributor_one_user_adjust_power_success(
     ) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works with single user works."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel1 = Bidirectional[Request, Result]("user1", "power_distributor")
         service_channels = {
@@ -461,9 +440,6 @@ async def test_power_distributor_one_user_adjust_power_success(
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor(service_channels)
 
@@ -492,7 +468,7 @@ async def test_power_distributor_one_user_adjust_power_success(
     async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> None:
         # pylint: disable=too-many-locals
         """Test if power distribution works if not all batteries are working."""
-        microgrid = await self.init_mock_microgrid()
+        await self.init_mock_microgrid(mocker)
 
         channel = Bidirectional[Request, Result]("user1", "power_distributor")
 
@@ -501,9 +477,6 @@ async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> Non
         )
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
-        mocker.patch(
-            "frequenz.sdk.microgrid.get", return_value=microgrid.mock_microgrid
-        )
 
         distributor = PowerDistributingActor({"user1": channel.service_handle})
 

tests/utils/mock_microgrid.py

@@ -8,6 +8,7 @@
 
 from frequenz.channels import Broadcast, Receiver
 from google.protobuf.empty_pb2 import Empty  # pylint: disable=no-name-in-module
+from pytest_mock import MockerFixture
 
 from frequenz.sdk.microgrid import Microgrid
 from frequenz.sdk.microgrid._graph import ComponentGraph, _MicrogridComponentGraph
@@ -55,6 +56,17 @@ def __init__(self, components: Set[Component], connections: Set[Connection]):
             id: channel.new_sender() for id, channel in inv_channels.items()
         }
 
+    async def initialize(self, mocker: MockerFixture) -> None:
+        """Mock `microgrid.get` call to return this mock_microgrid.
+
+        Args:
+            mocker: mocker from the current test
+        """
+        # Mock _MICROGRID, so `get` method return this mocked microgrid.
+        mocker.patch(
+            "frequenz.sdk.microgrid._microgrid._MICROGRID", self.mock_microgrid
+        )
+
     @property
     def mock_microgrid(self) -> Microgrid:
         """Return mock microgrid.
@@ -92,8 +104,10 @@ async def send(self, data: Union[BatteryData, InverterData]) -> bool:
         cid = data.component_id
         if isinstance(data, BatteryData):
             return await self._battery_data_senders[cid].send(data)
+        if isinstance(data, InverterData):
+            return await self._inverter_data_senders[cid].send(data)
 
-        return await self._inverter_data_senders[cid].send(data)
+        raise RuntimeError(f"{type(data)} is not supported in MockMicrogridClient.")
 
     def _create_battery_channels(self) -> Dict[int, Broadcast[BatteryData]]:
         """Create channels for the batteries.