stream() Return streamer object instead of receiver object

Signed-off-by: Mathias L. Baumann <mathias.baumann@frequenz.com>

Author: Marenz

src/frequenz/client/dispatch/_client.py

@@ -31,12 +31,11 @@
     UpdateMicrogridDispatchResponse,
 )
 
-from frequenz import channels
 from frequenz.client.base.channel import ChannelOptions, SslOptions
 from frequenz.client.base.client import BaseApiClient
 from frequenz.client.base.conversion import to_timestamp
 from frequenz.client.base.exception import ClientNotConnected
-from frequenz.client.base.retry import LinearBackoff
+from frequenz.client.base.retry import LinearBackoff, Strategy
 from frequenz.client.base.streaming import GrpcStreamBroadcaster
 
 from ._internal_types import DispatchCreateRequest
@@ -61,13 +60,16 @@ def __init__(
         server_url: str,
         key: str,
         connect: bool = True,
+        stream_retry_strategy: Strategy | None = None,
     ) -> None:
         """Initialize the client.
 
         Args:
             server_url: The URL of the server to connect to.
             key: API key to use for authentication.
-            connect: Whether to connect to the service immediately.
+            connect: Whether to connect to the service immediately (default: True).
+            stream_retry_strategy: The retry strategy to use for streaming
+                methods, defaults to `LinearBackoff` with 1 second interval and no limit.
         """
         super().__init__(
             server_url,
@@ -89,6 +91,11 @@ def __init__(
         self._streams: dict[
             int, GrpcStreamBroadcaster[StreamMicrogridDispatchesResponse, DispatchEvent]
         ] = {}
+        self._stream_retry_strategy = (
+            LinearBackoff(interval=1, limit=None)
+            if stream_retry_strategy is None
+            else stream_retry_strategy.copy()
+        )
         """A dictionary of streamers, keyed by microgrid_id."""
 
     @property
@@ -197,7 +204,9 @@ def to_interval(
             else:
                 break
 
-    def stream(self, microgrid_id: int) -> channels.Receiver[DispatchEvent]:
+    def stream(
+        self, microgrid_id: int
+    ) -> GrpcStreamBroadcaster[StreamMicrogridDispatchesResponse, DispatchEvent]:
         """Receive a stream of dispatch events.
 
         This function returns a receiver channel that can be used to receive
@@ -216,14 +225,8 @@ def stream(self, microgrid_id: int) -> channels.Receiver[DispatchEvent]:
             microgrid_id: The microgrid_id to receive dispatches for.
 
         Returns:
-            A receiver channel to receive the stream of dispatch events.
+            A broadcaster to receive the stream of dispatch events.
         """
-        return self._get_stream(microgrid_id).new_receiver()
-
-    def _get_stream(
-        self, microgrid_id: int
-    ) -> GrpcStreamBroadcaster[StreamMicrogridDispatchesResponse, DispatchEvent]:
-        """Get an instance to the streaming helper."""
         broadcaster = self._streams.get(microgrid_id)
         # pylint: disable=protected-access
         if broadcaster is not None and broadcaster._channel.is_closed:
@@ -241,7 +244,7 @@ def _get_stream(
                     ),
                 ),
                 transform=DispatchEvent.from_protobuf,
-                retry_strategy=LinearBackoff(interval=1, limit=None),
+                retry_strategy=self._stream_retry_strategy.copy(),
             )
             self._streams[microgrid_id] = broadcaster
 

tests/test_client.py

@@ -307,7 +307,7 @@ async def test_dispatch_stream(client: FakeClient, sample: Dispatch) -> None:
     microgrid_id = random.randint(1, 100)
     dispatches = [sample, sample, sample]
 
-    stream = client.stream(microgrid_id)
+    stream = client.stream(microgrid_id).new_receiver()
 
     async def expect(dispatch: Dispatch, event: Event) -> None:
         message = await stream.receive()