Return GrpcStreamBroadcaster instances from the streaming methods

This makes it easier to close the streamers when they are no longer
needed.

Also rename the methods from `stream_*` to `*_stream`.

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

Author: shsms

src/frequenz/client/electricity_trading/_client.py

@@ -24,7 +24,6 @@
 from frequenz.api.electricity_trading.v1.electricity_trading_pb2_grpc import (
     ElectricityTradingServiceStub,
 )
-from frequenz.channels import Receiver
 from frequenz.client.base.client import BaseApiClient
 from frequenz.client.base.exception import ClientNotConnected
 from frequenz.client.base.streaming import GrpcStreamBroadcaster
@@ -218,7 +217,7 @@ def stub(self) -> electricity_trading_pb2_grpc.ElectricityTradingServiceAsyncStu
         # type-checker, so it can only be used for type hints.
         return self._stub  # type: ignore
 
-    def stream_gridpool_orders(
+    def gridpool_orders_stream(
         # pylint: disable=too-many-arguments, too-many-positional-arguments
         self,
         gridpool_id: int,
@@ -227,9 +226,9 @@ def stream_gridpool_orders(
         delivery_area: DeliveryArea | None = None,
         delivery_period: DeliveryPeriod | None = None,
         tag: str | None = None,
-        max_size: int = 50,
-        warn_on_overflow: bool = False,
-    ) -> Receiver[OrderDetail]:
+    ) -> GrpcStreamBroadcaster[
+        electricity_trading_pb2.ReceiveGridpoolOrdersStreamResponse, OrderDetail
+    ]:
         """
         Stream gridpool orders.
 
@@ -240,10 +239,6 @@ def stream_gridpool_orders(
             delivery_area: Delivery area to filter for.
             delivery_period: Delivery period to filter for.
             tag: Tag to filter for.
-            max_size: The maximum number of messages to buffer.
-            warn_on_overflow: Whether to log a warning when the receiver's
-                buffer is full and a message is dropped.
-
 
         Returns:
             Async generator of orders.
@@ -281,11 +276,9 @@ def stream_gridpool_orders(
                     "Error occurred while streaming gridpool orders: %s", e
                 )
                 raise
-        return self._gridpool_orders_streams[stream_key].new_receiver(
-            warn_on_overflow=warn_on_overflow, maxsize=max_size
-        )
+        return self._gridpool_orders_streams[stream_key]
 
-    def stream_gridpool_trades(
+    def gridpool_trades_stream(
         # pylint: disable=too-many-arguments, too-many-positional-arguments
         self,
         gridpool_id: int,
@@ -294,9 +287,9 @@ def stream_gridpool_trades(
         market_side: MarketSide | None = None,
         delivery_period: DeliveryPeriod | None = None,
         delivery_area: DeliveryArea | None = None,
-        max_size: int = 50,
-        warn_on_overflow: bool = False,
-    ) -> Receiver[Trade]:
+    ) -> GrpcStreamBroadcaster[
+        electricity_trading_pb2.ReceiveGridpoolTradesStreamResponse, Trade
+    ]:
         """
         Stream gridpool trades.
 
@@ -307,9 +300,6 @@ def stream_gridpool_trades(
             market_side: The market side to filter for.
             delivery_period: The delivery period to filter for.
             delivery_area: The delivery area to filter for.
-            max_size: The maximum number of messages to buffer.
-            warn_on_overflow: Whether to log a warning when the receiver's
-                buffer is full and a message is dropped.
 
         Returns:
             The gridpool trades streamer.
@@ -347,20 +337,18 @@ def stream_gridpool_trades(
                     "Error occurred while streaming gridpool trades: %s", e
                 )
                 raise
-        return self._gridpool_trades_streams[stream_key].new_receiver(
-            warn_on_overflow=warn_on_overflow, maxsize=max_size
-        )
+        return self._gridpool_trades_streams[stream_key]
 
-    def stream_public_trades(
+    def public_trades_stream(
         # pylint: disable=too-many-arguments, too-many-positional-arguments
         self,
         states: list[TradeState] | None = None,
         delivery_period: DeliveryPeriod | None = None,
         buy_delivery_area: DeliveryArea | None = None,
         sell_delivery_area: DeliveryArea | None = None,
-        max_size: int = 50,
-        warn_on_overflow: bool = False,
-    ) -> Receiver[PublicTrade]:
+    ) -> GrpcStreamBroadcaster[
+        electricity_trading_pb2.ReceivePublicTradesStreamResponse, PublicTrade
+    ]:
         """
         Stream public trades.
 
@@ -369,9 +357,6 @@ def stream_public_trades(
             delivery_period: Delivery period to filter for.
             buy_delivery_area: Buy delivery area to filter for.
             sell_delivery_area: Sell delivery area to filter for.
-            max_size: The maximum number of messages to buffer.
-            warn_on_overflow: Whether to log a warning when the receiver's
-                buffer is full and a message is dropped.
 
         Returns:
             Async generator of orders.
@@ -405,9 +390,7 @@ def stream_public_trades(
             except grpc.RpcError as e:
                 _logger.exception("Error occurred while streaming public trades: %s", e)
                 raise
-        return self._public_trades_streams[public_trade_filter].new_receiver(
-            warn_on_overflow=warn_on_overflow, maxsize=max_size
-        )
+        return self._public_trades_streams[public_trade_filter]
 
     def validate_params(
         # pylint: disable=too-many-arguments, too-many-positional-arguments, too-many-branches

src/frequenz/client/electricity_trading/cli/etrading.py

@@ -73,7 +73,7 @@ async def list_public_trades(url: str, key: str, *, delivery_start: datetime) ->
         if delivery_start <= datetime.now(timezone.utc):
             return
 
-    stream = client.stream_public_trades(delivery_period=delivery_period)
+    stream = client.public_trades_stream(delivery_period=delivery_period).new_receiver()
     async for trade in stream:
         print_public_trade(trade)
 
@@ -111,7 +111,9 @@ async def list_gridpool_trades(
     if delivery_start and delivery_start <= datetime.now(timezone.utc):
         return
 
-    stream = client.stream_gridpool_trades(gid, delivery_period=delivery_period)
+    stream = client.gridpool_trades_stream(
+        gid, delivery_period=delivery_period
+    ).new_receiver()
     async for trade in stream:
         print_trade(trade)
 
@@ -154,7 +156,9 @@ async def list_gridpool_orders(
     if delivery_start and delivery_start <= datetime.now(timezone.utc):
         return
 
-    stream = client.stream_gridpool_orders(gid, delivery_period=delivery_period)
+    stream = client.gridpool_orders_stream(
+        gid, delivery_period=delivery_period
+    ).new_receiver()
     async for order in stream:
         print_order(order)
 

tests/test_client.py

@@ -133,7 +133,7 @@ def set_up_order_detail_response(
 
 async def test_stream_gridpool_orders(set_up: SetupParams) -> None:
     """Test the method streaming gridpool orders."""
-    set_up.client.stream_gridpool_orders(set_up.gridpool_id)
+    set_up.client.gridpool_orders_stream(set_up.gridpool_id)
     await asyncio.sleep(0)
 
     set_up.mock_stub.ReceiveGridpoolOrdersStream.assert_called_once()
@@ -146,7 +146,7 @@ async def test_stream_gridpool_orders_with_optional_inputs(set_up: SetupParams)
     # Fields to filter for
     order_states = [OrderState.ACTIVE]
 
-    set_up.client.stream_gridpool_orders(set_up.gridpool_id, order_states=order_states)
+    set_up.client.gridpool_orders_stream(set_up.gridpool_id, order_states=order_states)
     await asyncio.sleep(0)
 
     set_up.mock_stub.ReceiveGridpoolOrdersStream.assert_called_once()
@@ -161,7 +161,7 @@ async def test_stream_gridpool_trades(
     set_up: SetupParams,
 ) -> None:
     """Test the method streaming gridpool trades."""
-    set_up.client.stream_gridpool_trades(
+    set_up.client.gridpool_trades_stream(
         gridpool_id=set_up.gridpool_id, market_side=set_up.side
     )
     await asyncio.sleep(0)
@@ -179,7 +179,7 @@ async def test_stream_public_trades(
     # Fields to filter for
     trade_states = [TradeState.ACTIVE]
 
-    set_up.client.stream_public_trades(states=trade_states)
+    set_up.client.public_trades_stream(states=trade_states)
     await asyncio.sleep(0)
 
     set_up.mock_stub.ReceivePublicTradesStream.assert_called_once()