Support multiple metrics

Signed-off-by: cwasicki <[email protected]>
cwasicki · Mar 22, 2024 · 8a3f4d0 · 8a3f4d0
1 parent b125855
commit 8a3f4d0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/examples/client.py b/examples/client.py
@@ -98,8 +98,6 @@ async def run(
 
     metrics = [Metric[mn] for mn in metric_names]
 
-    assert len(metrics) == 1, "Only single metric is supported"
-
     def data_iter() -> AsyncIterator[MetricSample]:
         """Iterate over single metric.
 
@@ -108,10 +106,10 @@ def data_iter() -> AsyncIterator[MetricSample]:
         Returns:
             Iterator over single metric samples
         """
-        return client.iterate_single_metric(
+        return client.iterate_single_component(
             microgrid_id=microgrid_id,
             component_id=component_id,
-            metric=metrics[0],
+            metrics=metrics,
             start_dt=start_dt,
             end_dt=end_dt,
             page_size=page_size,

diff --git a/src/frequenz/client/reporting/_client.py b/src/frequenz/client/reporting/_client.py
@@ -117,12 +117,12 @@ def __init__(self, service_address: str):
         self._stub = ReportingStub(self._grpc_channel)
 
     # pylint: disable=too-many-arguments
-    async def iterate_single_metric(
+    async def iterate_single_component(
         self,
         *,
         microgrid_id: int,
         component_id: int,
-        metric: Metric,
+        metrics: Metric | list[Metric],
         start_dt: datetime,
         end_dt: datetime,
         page_size: int = 1000,
@@ -132,7 +132,7 @@ async def iterate_single_metric(
         Args:
             microgrid_id: The microgrid ID.
             component_id: The component ID.
-            metric: The metric name.
+            metrics: The metric name or list of metric names.
             start_dt: The start date and time.
             end_dt: The end date and time.
             page_size: The page size.
@@ -144,7 +144,7 @@ async def iterate_single_metric(
         """
         async for page in self._iterate_components_data_pages(
             microgrid_components=[(microgrid_id, [component_id])],
-            metrics=[metric],
+            metrics=[metrics] if isinstance(metrics, Metric) else metrics,
             start_dt=start_dt,
             end_dt=end_dt,
             page_size=page_size,