Implement ReceiveComponentDataStreamRequest

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

src/frequenz/client/microgrid/_client.py

@@ -17,6 +17,7 @@
 from frequenz.api.common.v1.metrics import bounds_pb2, metric_sample_pb2
 from frequenz.api.common.v1.microgrid.components import components_pb2
 from frequenz.api.microgrid.v1 import microgrid_pb2, microgrid_pb2_grpc
+from frequenz.channels import Receiver
 from frequenz.client.base import channel, client, conversion, retry, streaming
 from frequenz.client.common.microgrid.components import ComponentId
 from google.protobuf.empty_pb2 import Empty
@@ -30,6 +31,8 @@
 from .component._component_proto import component_from_proto
 from .component._connection import ComponentConnection
 from .component._connection_proto import component_connection_from_proto
+from .component._data_samples import ComponentDataSamples
+from .component._data_samples_proto import component_data_samples_from_proto
 from .component._types import ComponentTypes
 from .metrics._bounds import Bounds
 from .metrics._metric import Metric
@@ -84,8 +87,11 @@ def __init__(
             connect=connect,
             channel_defaults=channel_defaults,
         )
-        self._broadcasters: dict[
-            ComponentId, streaming.GrpcStreamBroadcaster[Any, Any]
+        self._component_data_broadcasters: dict[
+            str,
+            streaming.GrpcStreamBroadcaster[
+                microgrid_pb2.ReceiveComponentDataStreamResponse, ComponentDataSamples
+            ],
         ] = {}
         self._sensor_data_broadcasters: dict[
             str,
@@ -118,15 +124,15 @@ async def __aexit__(
                 *(
                     broadcaster.stop()
                     for broadcaster in itertools.chain(
-                        self._broadcasters.values(),
+                        self._component_data_broadcasters.values(),
                         self._sensor_data_broadcasters.values(),
                     )
                 ),
                 return_exceptions=True,
             )
             if isinstance(exc, BaseException)
         )
-        self._broadcasters.clear()
+        self._component_data_broadcasters.clear()
         self._sensor_data_broadcasters.clear()
 
         result = None
@@ -530,6 +536,70 @@ async def add_component_bounds(  # noqa: DOC502 (Raises ApiClientError indirectl
 
         return None
 
+    # noqa: DOC502 (Raises ApiClientError indirectly)
+    def receive_component_data_samples_stream(
+        self,
+        component: ComponentId | Component,
+        metrics: Iterable[Metric | int],
+        *,
+        buffer_size: int = 50,
+    ) -> Receiver[ComponentDataSamples]:
+        """Stream data samples from a component.
+
+        At least one metric must be specified. If no metric is specified, then the
+        stream will raise an error.
+
+        Warning:
+            Components may not support all metrics. If a component does not support
+            a given metric, then the returned data stream will not contain that metric.
+
+            There is no way to tell if a metric is not being received because the
+            component does not support it or because there is a transient issue when
+            retrieving the metric from the component.
+
+            The supported metrics by a component can even change with time, for example,
+            if a component is updated with new firmware.
+
+        Args:
+            component: The component to stream data from.
+            metrics: List of metrics to return. Only the specified metrics will be
+                returned.
+            buffer_size: The maximum number of messages to buffer in the returned
+                receiver. After this limit is reached, the oldest messages will be
+                dropped.
+
+        Returns:
+            The data stream from the component.
+        """
+        component_id = _get_component_id(component)
+        metrics_set = frozenset([_get_metric_value(m) for m in metrics])
+        key = f"{component_id}-{hash(metrics_set)}"
+        broadcaster = self._component_data_broadcasters.get(key)
+        if broadcaster is None:
+            client_id = hex(id(self))[2:]
+            stream_name = f"microgrid-client-{client_id}-component-data-{key}"
+            # Alias to avoid too long lines linter errors
+            # pylint: disable-next=invalid-name
+            Request = microgrid_pb2.ReceiveComponentDataStreamRequest
+            broadcaster = streaming.GrpcStreamBroadcaster(
+                stream_name,
+                lambda: aiter(
+                    self.stub.ReceiveComponentDataStream(
+                        Request(
+                            component_id=_get_component_id(component),
+                            filter=Request.ComponentDataStreamFilter(
+                                metrics=metrics_set
+                            ),
+                        ),
+                        timeout=DEFAULT_GRPC_CALL_TIMEOUT,
+                    )
+                ),
+                lambda msg: component_data_samples_from_proto(msg.data),
+                retry_strategy=self._retry_strategy,
+            )
+            self._component_data_broadcasters[key] = broadcaster
+        return broadcaster.new_receiver(maxsize=buffer_size)
+
 
 class Validity(enum.Enum):
     """The duration for which a given list of bounds will stay in effect."""

tests/client_test_cases/receive_component_data_samples_stream/empty_case.py

@@ -0,0 +1,48 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""Test data for empty component data stream."""
+
+from typing import Any
+
+import pytest
+from frequenz.api.common.v1.metrics import metric_sample_pb2
+from frequenz.api.common.v1.microgrid.components import components_pb2
+from frequenz.api.microgrid.v1 import microgrid_pb2
+from frequenz.channels import Receiver, ReceiverStoppedError
+from frequenz.client.common.microgrid.components import ComponentId
+
+from frequenz.client.microgrid.component import ComponentDataSamples
+
+client_args = (ComponentId(1), [metric_sample_pb2.Metric.METRIC_AC_CURRENT])
+
+
+def assert_stub_method_call(stub_method: Any) -> None:
+    """Assert that the gRPC request matches the expected request."""
+    stub_method.assert_called_once_with(
+        microgrid_pb2.ReceiveComponentDataStreamRequest(
+            component_id=1,
+            filter=microgrid_pb2.ReceiveComponentDataStreamRequest.ComponentDataStreamFilter(
+                metrics=[metric_sample_pb2.Metric.METRIC_AC_CURRENT]
+            ),
+        ),
+        timeout=60.0,
+    )
+
+
+# The mock response from the server
+grpc_response = microgrid_pb2.ReceiveComponentDataStreamResponse(
+    data=components_pb2.ComponentData(component_id=1, metric_samples=[], states=[]),
+)
+
+
+# The expected result from the client method
+async def assert_client_result(receiver: Receiver[Any]) -> None:
+    """Assert that the client result matches the expected empty data."""
+    result = await receiver.receive()
+    assert result == ComponentDataSamples(
+        component_id=ComponentId(1), metric_samples=[], states=[]
+    )
+
+    with pytest.raises(ReceiverStoppedError):
+        await receiver.receive()

tests/client_test_cases/receive_component_data_samples_stream/error_case.py

@@ -0,0 +1,85 @@
+# License: MIT
+# Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+"""Test data for component data stream with error."""
+
+import enum
+from collections.abc import AsyncIterator
+from typing import Any
+
+import pytest
+from frequenz.api.common.v1.metrics import metric_sample_pb2
+from frequenz.api.common.v1.microgrid.components import components_pb2
+from frequenz.api.microgrid.v1 import microgrid_pb2
+from frequenz.channels import Receiver, ReceiverStoppedError
+from frequenz.client.common.microgrid.components import ComponentId
+from grpc import StatusCode
+
+from frequenz.client.microgrid.component import ComponentDataSamples
+from tests.util import make_grpc_error
+
+client_args = (ComponentId(1), [metric_sample_pb2.Metric.METRIC_DC_VOLTAGE])
+
+
+def assert_stub_method_call(stub_method: Any) -> None:
+    """Assert that the gRPC request matches the expected request."""
+    stub_method.assert_called_once_with(
+        microgrid_pb2.ReceiveComponentDataStreamRequest(
+            component_id=1,
+            filter=microgrid_pb2.ReceiveComponentDataStreamRequest.ComponentDataStreamFilter(
+                metrics=[metric_sample_pb2.Metric.METRIC_DC_VOLTAGE]
+            ),
+        ),
+        timeout=60.0,
+    )
+
+
+@enum.unique
+class _State(enum.Enum):
+    """State of the gRPC response simulation."""
+
+    INITIAL = "initial"
+    ERROR = "error"
+    RECEIVING = "receiving"
+
+
+_iterations = 0
+_state: _State = _State.INITIAL
+
+
+async def grpc_response() -> AsyncIterator[Any]:
+    """Simulate a gRPC response with an error on the first iteration."""
+    global _iterations, _state  # pylint: disable=global-statement
+
+    _iterations += 1
+    if _iterations == 1:
+        _state = _State.ERROR
+        raise make_grpc_error(StatusCode.UNAVAILABLE)
+
+    _state = _State.RECEIVING
+    for _ in range(3):
+        yield microgrid_pb2.ReceiveComponentDataStreamResponse(
+            data=components_pb2.ComponentData(
+                component_id=1, metric_samples=[], states=[]
+            ),
+        )
+
+
+# The expected result from the client method (exception in this case)
+async def assert_client_result(receiver: Receiver[Any]) -> None:
+    """Assert that the client can keep receiving data after an error."""
+    assert _state is _State.ERROR
+
+    async for result in receiver:
+        assert result == ComponentDataSamples(
+            component_id=ComponentId(1), metric_samples=[], states=[]
+        )
+        # We need the type ignore here because mypy doesn't realize _state is
+        # global and updated from outside this function, so it wrongly narrows
+        # its type to `Literal[_State.ERROR]`, and complaining about the
+        # impossibility of overlapping with _STATE.RECEIVING.
+        # https://github.com/python/mypy/issues/19283
+        assert _state is _State.RECEIVING  # type: ignore[comparison-overlap]
+
+    with pytest.raises(ReceiverStoppedError):
+        await receiver.receive()