Add a stream_sensor_data() RPC to get sensor data

This RPC is also based on version 0.17 of the microgrid API instead of
the current 0.15 version, so users can already use the new interface
and make migrating to 0.17 more straight forward.

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

Author: llucax

src/frequenz/client/microgrid/_client.py

@@ -6,13 +6,15 @@
 from __future__ import annotations
 
 import asyncio
+import itertools
 import logging
 from collections.abc import Callable, Iterable, Set
 from dataclasses import replace
-from typing import Any, TypeVar
+from functools import partial
+from typing import Any, NotRequired, TypedDict, TypeVar, assert_never
 
 from frequenz.api.common import components_pb2, metrics_pb2
-from frequenz.api.microgrid import microgrid_pb2, microgrid_pb2_grpc
+from frequenz.api.microgrid import microgrid_pb2, microgrid_pb2_grpc, sensor_pb2
 from frequenz.channels import Receiver
 from frequenz.client.base import channel, client, retry, streaming
 from google.protobuf.empty_pb2 import Empty
@@ -35,10 +37,10 @@
 from ._connection import Connection
 from ._constants import RECEIVER_MAX_SIZE
 from ._exception import ApiClientError, ClientNotConnected
-from ._id import ComponentId, MicrogridId
+from ._id import ComponentId, MicrogridId, SensorId
 from ._metadata import Location, Metadata
-from ._sensor_proto import sensor_from_proto
-from .sensor import Sensor
+from ._sensor_proto import sensor_data_samples_from_proto, sensor_from_proto
+from .sensor import Sensor, SensorDataSamples, SensorMetric
 
 DEFAULT_GRPC_CALL_TIMEOUT = 60.0
 """The default timeout for gRPC calls made by this client (in seconds)."""
@@ -98,6 +100,12 @@ def __init__(
         self._broadcasters: dict[
             ComponentId, streaming.GrpcStreamBroadcaster[Any, Any]
         ] = {}
+        self._sensor_data_broadcasters: dict[
+            str,
+            streaming.GrpcStreamBroadcaster[
+                microgrid_pb2.ComponentData, SensorDataSamples
+            ],
+        ] = {}
         self._retry_strategy = retry_strategy
 
     @property
@@ -119,15 +127,22 @@ async def __aexit__(
         exc_tb: Any | None,
     ) -> bool | None:
         """Close the gRPC channel and stop all broadcasters."""
-        exceptions = [
+        exceptions = list(
             exc
             for exc in await asyncio.gather(
-                *(broadcaster.stop() for broadcaster in self._broadcasters.values()),
+                *(
+                    broadcaster.stop()
+                    for broadcaster in itertools.chain(
+                        self._broadcasters.values(),
+                        self._sensor_data_broadcasters.values(),
+                    )
+                ),
                 return_exceptions=True,
             )
             if isinstance(exc, BaseException)
-        ]
+        )
         self._broadcasters.clear()
+        self._sensor_data_broadcasters.clear()
 
         result = None
         try:
@@ -568,3 +583,91 @@ async def set_bounds(  # noqa: DOC503 (raises ApiClientError indirectly)
             ),
             method_name="AddInclusionBounds",
         )
+
+    # noqa: DOC502 (Raises ApiClientError indirectly)
+    def stream_sensor_data(
+        self,
+        sensor: SensorId | Sensor,
+        metrics: Iterable[SensorMetric | int] | None = None,
+        *,
+        buffer_size: int = 50,
+    ) -> Receiver[SensorDataSamples]:
+        """Stream data samples from a sensor.
+
+        Warning:
+            Sensors may not support all metrics. If a sensor 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
+            sensor does not support it or because there is a transient issue when
+            retrieving the metric from the sensor.
+
+            The supported metrics by a sensor can even change with time, for example,
+            if a sensor is updated with new firmware.
+
+        Args:
+            sensor: The sensor to stream data from.
+            metrics: If not `None`, only the specified metrics will be retrieved.
+                Otherwise all available metrics will be retrieved.
+            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:
+            A receiver to retrieve data from the sensor.
+        """
+        sensor_id = _get_sensor_id(sensor)
+        key = str(sensor_id)
+
+        class _ExtraArgs(TypedDict):
+            metrics: NotRequired[frozenset[sensor_pb2.SensorMetric.ValueType]]
+
+        extra_args: _ExtraArgs = {}
+        if metrics is not None:
+            extra_args["metrics"] = frozenset(
+                [_get_sensor_metric_value(m) for m in metrics]
+            )
+            # We use the frozenset because iterables are not hashable
+            key += f"{hash(extra_args['metrics'])}"
+
+        broadcaster = self._sensor_data_broadcasters.get(key)
+        if broadcaster is None:
+            client_id = hex(id(self))[2:]
+            stream_name = f"microgrid-client-{client_id}-sensor-data-{key}"
+            broadcaster = streaming.GrpcStreamBroadcaster(
+                stream_name,
+                lambda: aiter(
+                    self.stub.StreamComponentData(
+                        microgrid_pb2.ComponentIdParam(id=sensor_id),
+                        timeout=DEFAULT_GRPC_CALL_TIMEOUT,
+                    )
+                ),
+                partial(sensor_data_samples_from_proto, **extra_args),
+                retry_strategy=self._retry_strategy,
+            )
+            self._sensor_data_broadcasters[key] = broadcaster
+        return broadcaster.new_receiver(maxsize=buffer_size)
+
+
+def _get_sensor_id(sensor: SensorId | Sensor) -> int:
+    """Get the sensor ID from a sensor or sensor ID."""
+    match sensor:
+        case SensorId():
+            return int(sensor)
+        case Sensor():
+            return int(sensor.id)
+        case unexpected:
+            assert_never(unexpected)
+
+
+def _get_sensor_metric_value(
+    metric: SensorMetric | int,
+) -> sensor_pb2.SensorMetric.ValueType:
+    """Get the sensor metric ID from a sensor metric or sensor metric ID."""
+    match metric:
+        case SensorMetric():
+            return sensor_pb2.SensorMetric.ValueType(metric.value)
+        case int():
+            return sensor_pb2.SensorMetric.ValueType(metric)
+        case unexpected:
+            assert_never(unexpected)

tests/test_client.py

@@ -3,16 +3,20 @@
 
 """Tests for the microgrid client thin wrapper."""
 
+# We are going to split these tests in the future, but for now...
+# pylint: disable=too-many-lines
+
 import logging
 from collections.abc import AsyncIterator
+from datetime import datetime, timezone
 from typing import Any
 from unittest import mock
 
 import grpc.aio
 import pytest
 from frequenz.api.common import components_pb2, metrics_pb2
 from frequenz.api.microgrid import grid_pb2, inverter_pb2, microgrid_pb2, sensor_pb2
-from frequenz.client.base import retry
+from frequenz.client.base import conversion, retry
 from google.protobuf.empty_pb2 import Empty
 
 from frequenz.client.microgrid import (
@@ -33,7 +37,14 @@
     MicrogridId,
     SensorId,
 )
-from frequenz.client.microgrid.sensor import Sensor
+from frequenz.client.microgrid.sensor import (
+    Sensor,
+    SensorDataSamples,
+    SensorMetric,
+    SensorMetricSample,
+    SensorStateCode,
+    SensorStateSample,
+)
 
 
 class _TestClient(MicrogridApiClient):
@@ -590,15 +601,28 @@ def ev_charger101() -> microgrid_pb2.Component:
     )
 
 
+@pytest.fixture
+def sensor201() -> microgrid_pb2.Component:
+    """Return a test sensor component."""
+    return microgrid_pb2.Component(
+        id=201,
+        category=components_pb2.ComponentCategory.COMPONENT_CATEGORY_SENSOR,
+        sensor=sensor_pb2.Metadata(
+            type=components_pb2.SensorType.SENSOR_TYPE_THERMOMETER
+        ),
+    )
+
+
 @pytest.fixture
 def component_list(
     meter83: microgrid_pb2.Component,
     battery38: microgrid_pb2.Component,
     inverter99: microgrid_pb2.Component,
     ev_charger101: microgrid_pb2.Component,
+    sensor201: microgrid_pb2.Component,
 ) -> list[microgrid_pb2.Component]:
     """Return a list of test components."""
-    return [meter83, battery38, inverter99, ev_charger101]
+    return [meter83, battery38, inverter99, ev_charger101, sensor201]
 
 
 @pytest.mark.parametrize("method", ["meter_data", "battery_data", "inverter_data"])
@@ -890,6 +914,97 @@ async def test_set_bounds_grpc_error(client: _TestClient) -> None:
         await client.set_bounds(ComponentId(99), 0.0, 100.0)
 
 
+async def test_stream_sensor_data_success(
+    sensor201: microgrid_pb2.Component, client: _TestClient
+) -> None:
+    """Test successful streaming of sensor data."""
+    now = datetime.now(timezone.utc)
+
+    async def stream_data_impl(
+        *_: Any, **__: Any
+    ) -> AsyncIterator[microgrid_pb2.ComponentData]:
+        yield microgrid_pb2.ComponentData(
+            id=int(sensor201.id),
+            ts=conversion.to_timestamp(now),
+            sensor=sensor_pb2.Sensor(
+                state=sensor_pb2.State(
+                    component_state=sensor_pb2.ComponentState.COMPONENT_STATE_OK
+                ),
+                data=sensor_pb2.Data(
+                    sensor_data=[
+                        sensor_pb2.SensorData(
+                            value=1.0,
+                            sensor_metric=sensor_pb2.SensorMetric.SENSOR_METRIC_TEMPERATURE,
+                        )
+                    ],
+                ),
+            ),
+        )
+
+    client.mock_stub.StreamComponentData.side_effect = stream_data_impl
+    receiver = client.stream_sensor_data(
+        SensorId(sensor201.id), [SensorMetric.TEMPERATURE]
+    )
+    sample = await receiver.receive()
+
+    assert isinstance(sample, SensorDataSamples)
+    assert int(sample.sensor_id) == sensor201.id
+    assert sample.states == [
+        SensorStateSample(
+            sampled_at=now,
+            states=frozenset({SensorStateCode.ON}),
+            warnings=frozenset(),
+            errors=frozenset(),
+        )
+    ]
+    assert sample.metrics == [
+        SensorMetricSample(sampled_at=now, metric=SensorMetric.TEMPERATURE, value=1.0)
+    ]
+
+
+async def test_stream_sensor_data_grpc_error(
+    sensor201: microgrid_pb2.Component, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test stream_sensor_data() when the gRPC call fails and retries."""
+    caplog.set_level(logging.WARNING)
+
+    num_calls = 0
+
+    async def stream_data_error_impl(
+        *_: Any, **__: Any
+    ) -> AsyncIterator[microgrid_pb2.ComponentData]:
+        nonlocal num_calls
+        num_calls += 1
+        if num_calls <= 2:  # Fail first two times
+            raise grpc.aio.AioRpcError(
+                mock.MagicMock(name="mock_status"),
+                mock.MagicMock(name="mock_initial_metadata"),
+                mock.MagicMock(name="mock_trailing_metadata"),
+                f"fake grpc details stream_sensor_data num_calls={num_calls}",
+                "fake grpc debug_error_string",
+            )
+        # Succeed on the third call
+        yield microgrid_pb2.ComponentData(id=int(sensor201.id))
+
+    async with _TestClient(
+        retry_strategy=retry.LinearBackoff(interval=0.0, jitter=0.0, limit=3)
+    ) as client:
+        client.mock_stub.StreamComponentData.side_effect = stream_data_error_impl
+        receiver = client.stream_sensor_data(
+            SensorId(sensor201.id), [SensorMetric.TEMPERATURE]
+        )
+        sample = await receiver.receive()  # Should succeed after retries
+
+    assert isinstance(sample, SensorDataSamples)
+    assert int(sample.sensor_id) == sensor201.id
+
+    assert num_calls == 3  # Check that it was called 3 times (1 initial + 2 retries)
+    # Check log messages for retries
+    assert "connection ended, retrying" in caplog.text
+    assert "fake grpc details stream_sensor_data num_calls=1" in caplog.text
+    assert "fake grpc details stream_sensor_data num_calls=2" in caplog.text
+
+
 def _clear_components(component_list: microgrid_pb2.ComponentList) -> None:
     while component_list.components:
         component_list.components.pop()