Add Mock for initial_metadata() for base-clients GrpcStreamer

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

Author: Marenz

src/frequenz/client/dispatch/test/_service.py

@@ -5,10 +5,11 @@
 
 Useful for testing.
 """
+import asyncio
 import logging
 from dataclasses import dataclass, replace
 from datetime import datetime, timezone
-from typing import AsyncIterator
+from typing import AsyncIterator, TypeVar
 
 import grpc
 import grpc.aio
@@ -44,6 +45,33 @@
 _logger = logging.getLogger(__name__)
 
 
+T = TypeVar("T")
+
+
+class _MockStream(AsyncIterator[T]):
+    """A mock stream that wraps an async iterator and adds initial_metadata."""
+
+    def __init__(self, stream: AsyncIterator[T]) -> None:
+        """Initialize the mock stream.
+
+        Args:
+            stream: The stream to wrap.
+        """
+        self._iterator = stream.__aiter__()
+
+    async def initial_metadata(self) -> None:
+        """Do nothing, just to mock the grpc call."""
+        _logger.debug("Called initial_metadata()")
+
+    def __aiter__(self) -> AsyncIterator[T]:
+        """Return the async iterator."""
+        return self
+
+    async def __anext__(self) -> T:
+        """Return the next item from the stream."""
+        return await self._iterator.__anext__()
+
+
 class FakeService:
     """Dispatch mock service for testing."""
 
@@ -113,7 +141,7 @@ async def StreamMicrogridDispatches(
         self,
         request: StreamMicrogridDispatchesRequest,
         timeout: int = 5,  # pylint: disable=unused-argument
-    ) -> AsyncIterator[StreamMicrogridDispatchesResponse]:
+    ) -> _MockStream[StreamMicrogridDispatchesResponse]:
         """Stream microgrid dispatches changes.
 
         Args:
@@ -122,20 +150,28 @@ async def StreamMicrogridDispatches(
 
         Returns:
             An async generator for dispatch changes.
-
-        Yields:
-            An event for each dispatch change.
         """
-        receiver = self._stream_channel.new_receiver()
 
-        async for message in receiver:
-            _logger.debug("Received message: %s", message)
-            if message.microgrid_id == MicrogridId(request.microgrid_id):
-                response = StreamMicrogridDispatchesResponse(
-                    event=message.event.event.value,
-                    dispatch=message.event.dispatch.to_protobuf(),
-                )
-                yield response
+        async def stream() -> AsyncIterator[StreamMicrogridDispatchesResponse]:
+            """Stream microgrid dispatches changes."""
+            _logger.debug("Starting stream for microgrid %s", request.microgrid_id)
+            receiver = self._stream_channel.new_receiver()
+
+            async for message in receiver:
+                _logger.debug("Received message: %s", message)
+                if message.microgrid_id == MicrogridId(request.microgrid_id):
+                    response = StreamMicrogridDispatchesResponse(
+                        event=message.event.event.value,
+                        dispatch=message.event.dispatch.to_protobuf(),
+                    )
+                    yield response
+                else:
+                    _logger.debug(
+                        "Skipping message for microgrid %s",
+                        message.microgrid_id,
+                    )
+
+        return _MockStream(stream())
 
     # pylint: disable=too-many-branches
     @staticmethod
@@ -196,12 +232,18 @@ async def CreateMicrogridDispatch(
         # implicitly create the list if it doesn't exist
         self.dispatches.setdefault(microgrid_id, []).append(new_dispatch)
 
+        _logger.debug("Created new dispatch: %s", new_dispatch)
+
         await self._stream_sender.send(
             self.StreamEvent(
                 microgrid_id,
                 DispatchEvent(dispatch=new_dispatch, event=Event.CREATED),
             )
         )
+        # Give the stream a chance to process the message
+        await asyncio.sleep(0)
+
+        _logger.debug("Sent create event for dispatch: %s", new_dispatch)
 
         return CreateMicrogridDispatchResponse(dispatch=new_dispatch.to_protobuf())
 
@@ -293,6 +335,8 @@ async def UpdateMicrogridDispatch(
                 DispatchEvent(dispatch=dispatch, event=Event.UPDATED),
             )
         )
+        # Give the stream a chance to process the message
+        await asyncio.sleep(0)
 
         return UpdateMicrogridDispatchResponse(dispatch=dispatch.to_protobuf())
 
@@ -352,6 +396,8 @@ async def DeleteMicrogridDispatch(
                 ),
             )
         )
+        # Give the stream a chance to process the message
+        await asyncio.sleep(0)
 
         return Empty()