Update ComponentStatusTrackers to be BackgroundServices

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

Author: shsms

src/frequenz/sdk/actor/power_distributing/_component_pool_status_tracker.py

@@ -5,6 +5,7 @@
 
 
 import asyncio
+import contextlib
 import logging
 from collections import abc
 from datetime import timedelta
@@ -83,9 +84,6 @@ async def join(self) -> None:
     async def stop(self) -> None:
         """Stop the ComponentPoolStatusTracker instance."""
         await cancel_and_await(self._task)
-        await asyncio.gather(
-            *[tracker.stop() for tracker in self._component_status_trackers],
-        )
         await self._merged_status_receiver.stop()
 
     def _make_merged_status_receiver(
@@ -110,14 +108,17 @@ def _make_merged_status_receiver(
 
     async def _run(self) -> None:
         """Start tracking component status."""
-        while True:
-            try:
-                await self._update_status()
-            except Exception as err:  # pylint: disable=broad-except
-                _logger.error(
-                    "ComponentPoolStatus failed with error: %s. Restarting.", err
-                )
-                await asyncio.sleep(1.0)
+        async with contextlib.AsyncExitStack() as stack:
+            for tracker in self._component_status_trackers:
+                await stack.enter_async_context(tracker)
+            while True:
+                try:
+                    await self._update_status()
+                except Exception as err:  # pylint: disable=broad-except
+                    _logger.error(
+                        "ComponentPoolStatus failed with error: %s. Restarting.", err
+                    )
+                    await asyncio.sleep(1.0)
 
     async def _update_status(self) -> None:
         async for status in self._merged_status_receiver:

src/frequenz/sdk/actor/power_distributing/_component_status/_battery_status_tracker.py

@@ -20,14 +20,14 @@
 from frequenz.channels.util import Timer, select, selected_from
 from typing_extensions import override
 
-from ...._internal._asyncio import cancel_and_await
 from ....microgrid import connection_manager
 from ....microgrid.component import (
     BatteryData,
     ComponentCategory,
     ComponentData,
     InverterData,
 )
+from ..._background_service import BackgroundService
 from ._component_status import (
     ComponentStatus,
     ComponentStatusEnum,
@@ -126,7 +126,7 @@ def is_blocked(self) -> bool:
         return self.blocked_until > datetime.now(tz=timezone.utc)
 
 
-class BatteryStatusTracker(ComponentStatusTracker):
+class BatteryStatusTracker(ComponentStatusTracker, BackgroundService):
     """Class for tracking if battery is working.
 
     Status updates are sent out only when there is a status change.
@@ -186,7 +186,11 @@ def __init__(  # pylint: disable=too-many-arguments
         Raises:
             RuntimeError: If battery has no adjacent inverter.
         """
+        BackgroundService.__init__(self, name=f"BatteryStatusTracker({component_id})")
         self._max_data_age = max_data_age
+        self._status_sender = status_sender
+        self._set_power_result_receiver = set_power_result_receiver
+
         # First battery is considered as not working.
         # Change status after first messages are received.
         self._last_status: ComponentStatusEnum = ComponentStatusEnum.NOT_WORKING
@@ -212,8 +216,13 @@ def __init__(  # pylint: disable=too-many-arguments
 
         # Select needs receivers that can be get in async way only.
 
-        self._task: asyncio.Task[None] = asyncio.create_task(
-            self._run(status_sender, set_power_result_receiver)
+    @override
+    def start(self) -> None:
+        """Start the BatteryStatusTracker instance."""
+        self._tasks.add(
+            asyncio.create_task(
+                self._run(self._status_sender, self._set_power_result_receiver)
+            )
         )
 
     @property
@@ -225,10 +234,6 @@ def battery_id(self) -> int:
         """
         return self._battery.component_id
 
-    async def stop(self) -> None:
-        """Stop tracking battery status."""
-        await cancel_and_await(self._task)
-
     def _handle_status_battery(self, bat_data: BatteryData) -> None:
         self._battery.last_msg_correct = (
             self._is_message_reliable(bat_data)

src/frequenz/sdk/actor/power_distributing/_component_status/_component_status.py

@@ -12,6 +12,8 @@
 
 from frequenz.channels import Receiver, Sender
 
+from ..._background_service import BackgroundService
+
 
 @dataclass
 class ComponentPoolStatus:
@@ -77,11 +79,11 @@ class SetPowerResult:
     """Component IDs for which the last set power command failed."""
 
 
-class ComponentStatusTracker(ABC):
+class ComponentStatusTracker(BackgroundService, ABC):
     """Interface for specialized component status trackers to implement."""
 
     @abstractmethod
-    def __init__(  # pylint: disable=too-many-arguments
+    def __init__(  # pylint: disable=too-many-arguments,super-init-not-called
         self,
         component_id: int,
         max_data_age: timedelta,
@@ -102,7 +104,3 @@ def __init__(  # pylint: disable=too-many-arguments
             set_power_result_receiver: Channel to receive results of the requests to the
                 components.
         """
-
-    @abstractmethod
-    async def stop(self) -> None:
-        """Stop the ComponentStatusTracker instance."""

tests/actor/test_battery_status.py

@@ -7,10 +7,9 @@
 import asyncio
 import math
 from collections.abc import AsyncIterator, Iterable
-from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, Generic, TypeVar
+from typing import Generic, TypeVar
 
 import pytest
 
@@ -144,26 +143,6 @@ async def recv_timeout(recv: Receiver[T], timeout: float = 0.1) -> T | type[_Tim
         return _Timeout
 
 
-@asynccontextmanager
-async def battery_status_tracker(
-    *args: Any, **kwargs: Any
-) -> AsyncIterator[BatteryStatusTracker]:
-    """Create BatteryStatusTracker with given arguments.
-
-    Args:
-        *args: Arguments for BatteryStatusTracker.
-        **kwargs: Arguments for BatteryStatusTracker.
-
-    Yields:
-        BatteryStatusTracker with given arguments.
-    """
-    tracker = BatteryStatusTracker(*args, **kwargs)
-    try:
-        yield tracker
-    finally:
-        await tracker.stop()
-
-
 # pylint: disable=protected-access, unused-argument
 class TestBatteryStatus:
     """Tests BatteryStatusTracker."""
@@ -186,7 +165,7 @@ async def test_sync_update_status_with_messages(
         status_channel = Broadcast[ComponentStatus]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=5),
             max_blocking_duration=timedelta(seconds=30),
@@ -354,7 +333,7 @@ async def test_sync_blocking_feature(self, mocker: MockerFixture) -> None:
         status_channel = Broadcast[ComponentStatus]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             # increase max_data_age_sec for blocking tests.
             # Otherwise it will block blocking.
             BATTERY_ID,
@@ -492,7 +471,7 @@ async def test_sync_blocking_interrupted_with_with_max_data(
         status_channel = Broadcast[ComponentStatus]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=5),
             max_blocking_duration=timedelta(seconds=30),
@@ -540,7 +519,7 @@ async def test_sync_blocking_interrupted_with_invalid_message(
         status_channel = Broadcast[ComponentStatus]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=5),
             max_blocking_duration=timedelta(seconds=30),
@@ -601,7 +580,7 @@ async def test_timers(
         status_channel = Broadcast[ComponentStatus]("battery_status")
         set_power_result_channel = Broadcast[SetPowerResult]("set_power_result")
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=5),
             max_blocking_duration=timedelta(seconds=30),
@@ -667,7 +646,7 @@ async def test_async_battery_status(self, mocker: MockerFixture) -> None:
         status_receiver = status_channel.new_receiver()
         set_power_result_sender = set_power_result_channel.new_sender()
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=5),
             max_blocking_duration=timedelta(seconds=30),
@@ -751,7 +730,7 @@ async def setup_tracker(
 
         status_receiver = status_channel.new_receiver()
 
-        async with mock_microgrid, battery_status_tracker(
+        async with mock_microgrid, BatteryStatusTracker(
             BATTERY_ID,
             max_data_age=timedelta(seconds=0.1),
             max_blocking_duration=timedelta(seconds=1),