Add mocks for meter_data, ev_charger_data in MockMicrogridClient

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

Author: shsms

tests/utils/mock_microgrid.py

@@ -18,7 +18,9 @@
     Component,
     ComponentCategory,
     ComponentData,
+    EVChargerData,
     InverterData,
+    MeterData,
 )
 
 
@@ -42,8 +44,12 @@ def __init__(self, components: Set[Component], connections: Set[Connection]):
 
         bat_channels = self._create_battery_channels()
         inv_channels = self._create_inverter_channels()
+        meter_channels = self._create_meter_channels()
+        ev_charger_channels = self._create_ev_charger_channels()
 
-        mock_api = self._create_mock_api(components, bat_channels, inv_channels)
+        mock_api = self._create_mock_api(
+            components, bat_channels, inv_channels, meter_channels, ev_charger_channels
+        )
         kwargs: Dict[str, Any] = {
             "api_client": mock_api,
             "component_graph": self._component_graph,
@@ -56,6 +62,12 @@ def __init__(self, components: Set[Component], connections: Set[Connection]):
         self._inverter_data_senders = {
             id: channel.new_sender() for id, channel in inv_channels.items()
         }
+        self._meter_data_senders = {
+            id: channel.new_sender() for id, channel in meter_channels.items()
+        }
+        self._ev_charger_data_senders = {
+            id: channel.new_sender() for id, channel in ev_charger_channels.items()
+        }
 
     def initialize(self, mocker: MockerFixture) -> None:
         """Mock `microgrid.get` call to return this mock_microgrid.
@@ -107,15 +119,16 @@ async def send(self, data: ComponentData) -> bool:
             return await self._battery_data_senders[cid].send(data)
         if isinstance(data, InverterData):
             return await self._inverter_data_senders[cid].send(data)
+        if isinstance(data, MeterData):
+            return await self._meter_data_senders[cid].send(data)
+        if isinstance(data, EVChargerData):
+            return await self._ev_charger_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.
 
-        Args:
-            components: set of components.
-
         Returns:
             Dictionary where the key is battery id and the value is channel for this
                 battery.
@@ -131,15 +144,28 @@ def _create_battery_channels(self) -> Dict[int, Broadcast[BatteryData]]:
             bid: Broadcast[BatteryData]("battery_data_" + str(bid)) for bid in batteries
         }
 
+    def _create_meter_channels(self) -> Dict[int, Broadcast[MeterData]]:
+        """Create channels for the meters.
+
+        Returns:
+            Dictionary where the key is meter id and the value is channel for this
+                meter.
+        """
+        meters = [
+            c.component_id
+            for c in self.component_graph.components(
+                component_category={ComponentCategory.METER}
+            )
+        ]
+
+        return {cid: Broadcast[MeterData]("meter_data_" + str(cid)) for cid in meters}
+
     def _create_inverter_channels(self) -> Dict[int, Broadcast[InverterData]]:
         """Create channels for the inverters.
 
-        Args:
-            components: set of components.
-
         Returns:
-            Dictionary where the key is inverter id and the value is channel for this
-                battery.
+            Dictionary where the key is inverter id and the value is channel for
+                this inverter.
         """
         inverters = [
             c.component_id
@@ -153,11 +179,31 @@ def _create_inverter_channels(self) -> Dict[int, Broadcast[InverterData]]:
             for cid in inverters
         }
 
+    def _create_ev_charger_channels(self) -> Dict[int, Broadcast[EVChargerData]]:
+        """Create channels for the ev chargers.
+
+        Returns:
+            Dictionary where the key is the id of the ev_charger and the value is
+                channel for this ev_charger.
+        """
+        meters = [
+            c.component_id
+            for c in self.component_graph.components(
+                component_category={ComponentCategory.EV_CHARGER}
+            )
+        ]
+
+        return {
+            cid: Broadcast[EVChargerData]("meter_data_" + str(cid)) for cid in meters
+        }
+
     def _create_mock_api(
         self,
         components: Set[Component],
-        bat_channel: Dict[int, Broadcast[BatteryData]],
-        inv_channel: Dict[int, Broadcast[InverterData]],
+        bat_channels: Dict[int, Broadcast[BatteryData]],
+        inv_channels: Dict[int, Broadcast[InverterData]],
+        meter_channels: Dict[int, Broadcast[MeterData]],
+        ev_charger_channels: Dict[int, Broadcast[EVChargerData]],
     ) -> MagicMock:
         """Create mock of MicrogridApiClient.
 
@@ -176,11 +222,21 @@ def _create_mock_api(
         # NOTE that has to be partial, because battery_data has id argument and takes
         # channel based on the argument.
         api.battery_data = AsyncMock(
-            side_effect=partial(self._get_battery_receiver, channels=bat_channel)
+            side_effect=partial(self._get_battery_receiver, channels=bat_channels)
         )
 
         api.inverter_data = AsyncMock(
-            side_effect=partial(self._get_inverter_receiver, channels=inv_channel)
+            side_effect=partial(self._get_inverter_receiver, channels=inv_channels)
+        )
+
+        api.meter_data = AsyncMock(
+            side_effect=partial(self._get_meter_receiver, channels=meter_channels)
+        )
+
+        api.ev_charger_data = AsyncMock(
+            side_effect=partial(
+                self._get_ev_charger_receiver, channels=ev_charger_channels
+            )
         )
 
         # Can be override in the future
@@ -199,8 +255,39 @@ def _get_battery_receiver(
             channels: Broadcast channels
 
         Returns:
-            Receiver[Union[BatteryData, InverterData]]: Receiver from the given
-                channels.
+            Receiver from the given channels.
+        """
+        return channels[component_id].new_receiver("component" + str(component_id))
+
+    def _get_meter_receiver(
+        self,
+        component_id: int,
+        channels: Dict[int, Broadcast[MeterData]],
+    ) -> Receiver[MeterData]:
+        """Return receiver of the broadcast channel for given component_id.
+
+        Args:
+            component_id: component_id
+            channels: Broadcast channels
+
+        Returns:
+            Receiver from the given channels.
+        """
+        return channels[component_id].new_receiver("component" + str(component_id))
+
+    def _get_ev_charger_receiver(
+        self,
+        component_id: int,
+        channels: Dict[int, Broadcast[EVChargerData]],
+    ) -> Receiver[EVChargerData]:
+        """Return receiver of the broadcast channel for given component_id.
+
+        Args:
+            component_id: component_id
+            channels: Broadcast channels
+
+        Returns:
+            Receiver from the given channels.
         """
         return channels[component_id].new_receiver("component" + str(component_id))
 
@@ -216,7 +303,6 @@ def _get_inverter_receiver(
             channels: Broadcast channels
 
         Returns:
-            Receiver[Union[BatteryData, InverterData]]: Receiver from the given
-                channels.
+            Receiver from the given channels.
         """
         return channels[component_id].new_receiver("component" + str(component_id))