MockMigrogrid/Resampler: Fail on unhandled exceptions

The tasks created by the MockMigrogrid/Resampler were not handling
exceptions, so if anything goes wrong when streaming (or when setting up
the streaming), the task will silently die, which means no values are
sent to the channels, so tests will probably hang forever waiting for
values.

We now make sure that tasks are cleaned up from the lists when they are
done and also to exit the tests if there was an unhandled exception, so
at least the test will fail and show the origin of the raised exception.

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

Author: llucax

tests/timeseries/mock_microgrid.py

@@ -155,8 +155,18 @@ def inverters(comp_type: InverterType) -> list[int]:
         self.evc_component_states: dict[int, EVChargerComponentState] = {}
         self.evc_cable_states: dict[int, EVChargerCableState] = {}
 
-        self._streaming_coros: list[Coroutine[None, None, None]] = []
-        self._streaming_tasks: list[asyncio.Task[None]] = []
+        self._streaming_coros: list[tuple[int, Coroutine[None, None, None]]] = []
+        """The streaming coroutines for each component.
+
+        The tuple stores the component id we are streaming for as the first item and the
+        coroutine as the second item.
+        """
+
+        self._streaming_tasks: dict[int, asyncio.Task[None]] = {}
+        """The streaming tasks for each component.
+
+        The key is the component id we are streaming for in this task.
+        """
 
         if grid_meter:
             self._connect_to = self._grid_meter_id
@@ -218,9 +228,16 @@ def start_mock_client(
             A MockMicrogridClient instance.
         """
         self.init_mock_client(initialize_cb)
-        self._streaming_tasks = [
-            asyncio.create_task(coro) for coro in self._streaming_coros
-        ]
+
+        def _done_callback(task: asyncio.Task[None]) -> None:
+            if exc := task.exception():
+                raise SystemExit(f"Streaming task {task.get_name()!r} failed: {exc}")
+
+        for component_id, coro in self._streaming_coros:
+            task = asyncio.create_task(coro, name=f"component-id:{component_id}")
+            self._streaming_tasks[component_id] = task
+            task.add_done_callback(_done_callback)
+
         return self.mock_client
 
     async def _comp_data_send_task(
@@ -243,14 +260,17 @@ def _start_meter_streaming(self, meter_id: int) -> None:
         if not self._api_client_streaming:
             return
         self._streaming_coros.append(
-            self._comp_data_send_task(
+            (
                 meter_id,
-                lambda value, ts: MeterDataWrapper(
-                    component_id=meter_id,
-                    timestamp=ts,
-                    active_power=value,
-                    current_per_phase=(value + 100.0, value + 101.0, value + 102.0),
-                    voltage_per_phase=(value + 200.0, value + 199.8, value + 200.2),
+                self._comp_data_send_task(
+                    meter_id,
+                    lambda value, ts: MeterDataWrapper(
+                        component_id=meter_id,
+                        timestamp=ts,
+                        active_power=value,
+                        current_per_phase=(value + 100.0, value + 101.0, value + 102.0),
+                        voltage_per_phase=(value + 200.0, value + 199.8, value + 200.2),
+                    ),
                 ),
             )
         )
@@ -259,10 +279,13 @@ def _start_battery_streaming(self, bat_id: int) -> None:
         if not self._api_client_streaming:
             return
         self._streaming_coros.append(
-            self._comp_data_send_task(
+            (
                 bat_id,
-                lambda value, ts: BatteryDataWrapper(
-                    component_id=bat_id, timestamp=ts, soc=value
+                self._comp_data_send_task(
+                    bat_id,
+                    lambda value, ts: BatteryDataWrapper(
+                        component_id=bat_id, timestamp=ts, soc=value
+                    ),
                 ),
             )
         )
@@ -271,10 +294,13 @@ def _start_inverter_streaming(self, inv_id: int) -> None:
         if not self._api_client_streaming:
             return
         self._streaming_coros.append(
-            self._comp_data_send_task(
+            (
                 inv_id,
-                lambda value, ts: InverterDataWrapper(
-                    component_id=inv_id, timestamp=ts, active_power=value
+                self._comp_data_send_task(
+                    inv_id,
+                    lambda value, ts: InverterDataWrapper(
+                        component_id=inv_id, timestamp=ts, active_power=value
+                    ),
                 ),
             )
         )
@@ -283,17 +309,20 @@ def _start_ev_charger_streaming(self, evc_id: int) -> None:
         if not self._api_client_streaming:
             return
         self._streaming_coros.append(
-            self._comp_data_send_task(
+            (
                 evc_id,
-                lambda value, ts: EvChargerDataWrapper(
-                    component_id=evc_id,
-                    timestamp=ts,
-                    active_power=value,
-                    current_per_phase=(value + 10.0, value + 11.0, value + 12.0),
-                    component_state=self.evc_component_states[evc_id],
-                    cable_state=self.evc_cable_states[evc_id],
+                self._comp_data_send_task(
+                    evc_id,
+                    lambda value, ts: EvChargerDataWrapper(
+                        component_id=evc_id,
+                        timestamp=ts,
+                        active_power=value,
+                        current_per_phase=(value + 10.0, value + 11.0, value + 12.0),
+                        component_state=self.evc_component_states[evc_id],
+                        cable_state=self.evc_cable_states[evc_id],
+                    ),
                 ),
-            ),
+            )
         )
 
     def add_consumer_meters(self, count: int = 1) -> None:
@@ -563,10 +592,10 @@ async def cleanup(self) -> None:
 
         await self.mock_resampler._stop()
 
-        for coro in self._streaming_coros:
+        for _, coro in self._streaming_coros:
             coro.close()
 
-        for task in self._streaming_tasks:
+        for task in self._streaming_tasks.values():
             await cancel_and_await(task)
         microgrid.connection_manager._CONNECTION_MANAGER = None
         # pylint: enable=protected-access

tests/timeseries/mock_resampler.py

@@ -179,9 +179,13 @@ def voltage_senders(ids: list[int]) -> list[list[Sender[Sample[Quantity]]]]:
         )
 
         self._forward_tasks: dict[str, asyncio.Task[None]] = {}
-        self._request_handler_task = asyncio.create_task(
-            self._handle_resampling_requests()
-        )
+        task = asyncio.create_task(self._handle_resampling_requests())
+        task.add_done_callback(self._handle_task_done)
+        self._request_handler_task = task
+
+    def _handle_task_done(self, task: asyncio.Task[None]) -> None:
+        if exc := task.exception():
+            raise SystemExit(f"Task {task.get_name()!r} failed: {exc}") from exc
 
     async def _stop(self) -> None:
         tasks_to_stop = [
@@ -202,20 +206,29 @@ async def _channel_forward_messages(
 
     async def _handle_resampling_requests(self) -> None:
         async for request in self._resampler_request_channel.new_receiver():
-            if request.get_channel_name() in self._forward_tasks:
+            name = request.get_channel_name()
+            if name in self._forward_tasks:
                 continue
             basic_recv_name = f"{request.component_id}:{request.metric_id}"
             recv = self._basic_receivers[basic_recv_name].pop()
             assert recv is not None
-            self._forward_tasks[request.get_channel_name()] = asyncio.create_task(
+            task = asyncio.create_task(
                 self._channel_forward_messages(
                     recv,
                     self._channel_registry.get_or_create(
-                        Sample[Quantity], request.get_channel_name()
+                        Sample[Quantity], name
                     ).new_sender(),
-                )
+                ),
+                name=name,
             )
 
+            def _done_callback(task: asyncio.Task[None]) -> None:
+                del self._forward_tasks[task.get_name()]
+                self._handle_task_done(task)
+
+            task.add_done_callback(_done_callback)
+            self._forward_tasks[name] = task
+
     def make_sample(self, value: float | None) -> Sample[Quantity]:
         """Create a sample with the given value."""
         return Sample(