Update mock_microgrid to close channels when done sending.

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

tests/timeseries/mock_microgrid.py

@@ -148,6 +148,8 @@ async def _comp_data_send_task(
                 await self._microgrid.send(make_comp_data(val_to_send, timestamp))
             await asyncio.sleep(self._sample_rate_s)
 
+        await self._microgrid.close_channel(comp_id)
+
     def _start_meter_streaming(self, meter_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(

tests/utils/mock_microgrid.py

@@ -49,6 +49,13 @@ def __init__(self, components: Set[Component], connections: Set[Connection]):
         meter_channels = self._create_meter_channels()
         ev_charger_channels = self._create_ev_charger_channels()
 
+        self._all_channels: Dict[int, Broadcast[Any]] = {
+            **bat_channels,
+            **inv_channels,
+            **meter_channels,
+            **ev_charger_channels,
+        }
+
         mock_api = self._create_mock_api(
             bat_channels, inv_channels, meter_channels, ev_charger_channels
         )
@@ -128,6 +135,15 @@ async def send(self, data: ComponentData) -> bool:
 
         raise RuntimeError(f"{type(data)} is not supported in MockMicrogridClient.")
 
+    async def close_channel(self, cid: int) -> None:
+        """Close channel for given component id.
+
+        Args:
+            cid: Component id
+        """
+        if cid in self._all_channels:
+            await self._all_channels[cid].close()
+
     def _create_battery_channels(self) -> Dict[int, Broadcast[BatteryData]]:
         """Create channels for the batteries.