with now waits till dispatch is initialized

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

Author: Marenz

RELEASE_NOTES.md

@@ -25,3 +25,5 @@ This release introduces a more flexible and powerful mechanism for managing disp
     * `Dispatcher.start_dispatching(dispatch_type, actor_factory, merge_strategy)` to manage your actor for the given type and merge strategy. All you need provide is an actor factory.
     * `Dispatcher.stop_dispatching(dispatch_type)` to stop dispatching for the given type.
     * `Dispatcher.is_managed(dispatch_type)` to check if dispatching is active for the given type.
+* A new method `Dispatcher.wait_for_initialization()` has been added to wait for all actors to be initialized.
+* When using `async with Dispatcher(..) as dispatcher`, the dispatcher will first wait for the dispatch service to be initialized before entering the block.

src/frequenz/dispatch/_bg_service.py

@@ -129,6 +129,13 @@ def __init__(
         always at index 0.
         """
 
+        self._initial_fetch_event = asyncio.Event()
+        """The initial fetch event."""
+
+    async def wait_for_initialization(self) -> None:
+        """Wait for the initial fetch to complete."""
+        await self._initial_fetch_event.wait()
+
     # pylint: disable=redefined-builtin
     def new_lifecycle_events_receiver(self, type: str) -> Receiver[DispatchEvent]:
         """Create a new receiver for lifecycle events.
@@ -144,7 +151,10 @@ def new_lifecycle_events_receiver(self, type: str) -> Receiver[DispatchEvent]:
         )
 
     async def new_running_state_event_receiver(
-        self, type: str, *, merge_strategy: MergeStrategy | None = None
+        self,
+        type: str,
+        *,
+        merge_strategy: MergeStrategy | None = None,
     ) -> Receiver[Dispatch]:
         """Create a new receiver for running state events of the specified type.
 
@@ -171,6 +181,7 @@ async def new_running_state_event_receiver(
         Args:
             type: The type of events to receive.
             merge_strategy: The merge strategy to use.
+
         Returns:
             A new receiver for running state status.
         """
@@ -272,6 +283,8 @@ async def _fetch(self) -> None:
         This is used for the initial fetch and for re-fetching all dispatches
         if the connection was lost.
         """
+        self._initial_fetch_event.clear()
+
         old_dispatches = self._dispatches
         self._dispatches = {}
 
@@ -311,6 +324,8 @@ async def _fetch(self) -> None:
             dispatch._set_deleted()  # pylint: disable=protected-access
             await self._lifecycle_events_tx.send(Deleted(dispatch=dispatch))
 
+        self._initial_fetch_event.set()
+
     async def _update_dispatch_schedule_and_notify(
         self, dispatch: Dispatch | None, old_dispatch: Dispatch | None
     ) -> None:

src/frequenz/dispatch/_dispatcher.py

@@ -8,7 +8,7 @@
 import asyncio
 import logging
 from asyncio import Event
-from typing import Awaitable, Callable
+from typing import Awaitable, Callable, Self
 
 from frequenz.channels import Receiver
 from frequenz.client.dispatch import Client
@@ -221,6 +221,10 @@ def cancel(self, msg: str | None = None) -> None:
         for instance in self._actor_dispatchers.values():
             instance.cancel()
 
+    async def wait_for_initialization(self) -> None:
+        """Wait until the background service is initialized."""
+        await self._bg_service.wait_for_initialization()
+
     def is_managed(self, dispatch_type: str) -> bool:
         """Check if the dispatcher is managing actors for a given dispatch type.
 
@@ -271,7 +275,8 @@ def id_identity(dispatch: Dispatch) -> int:
         dispatcher = ActorDispatcher(
             actor_factory=actor_factory,
             running_status_receiver=await self.new_running_state_event_receiver(
-                dispatch_type, merge_strategy=merge_strategy
+                dispatch_type,
+                merge_strategy=merge_strategy,
             ),
             dispatch_identity=(
                 id_identity if merge_strategy is None else merge_strategy.identity
@@ -299,6 +304,19 @@ def client(self) -> Client:
         """Return the client."""
         return self._client
 
+    @override
+    async def __aenter__(self) -> Self:
+        """Enter an async context.
+
+        Start this background service.
+
+        Returns:
+            This background service.
+        """
+        self.start()
+        await self.wait_for_initialization()
+        return self
+
     def new_lifecycle_events_receiver(
         self, dispatch_type: str
     ) -> Receiver[DispatchEvent]:
@@ -368,5 +386,6 @@ async def new_running_state_event_receiver(
             A new receiver for dispatches whose running status changed.
         """
         return await self._bg_service.new_running_state_event_receiver(
-            dispatch_type, merge_strategy=merge_strategy
+            dispatch_type,
+            merge_strategy=merge_strategy,
         )

tests/test_mananging_actor.py

@@ -311,7 +311,9 @@ async def new_mock_receiver(
         # pylint: disable=protected-access
         assert "MANAGE_TEST" in dispatcher._actor_dispatchers
         actor_manager = dispatcher._actor_dispatchers["MANAGE_TEST"]
+        # pylint: disable=comparison-with-callable
         assert actor_manager._actor_factory == MockActor.create
+        # pylint: enable=comparison-with-callable
 
         dispatch = Dispatch(
             replace(