From 8a3f4d0eaaa599298ad8156c3f5a2bfff38f485e Mon Sep 17 00:00:00 2001
From: cwasicki <126617870+cwasicki@users.noreply.github.com>
Date: Fri, 22 Mar 2024 17:52:59 +0100
Subject: [PATCH] Support multiple metrics

Signed-off-by: cwasicki <126617870+cwasicki@users.noreply.github.com>
---
 examples/client.py                       | 6 ++----
 src/frequenz/client/reporting/_client.py | 8 ++++----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/examples/client.py b/examples/client.py
index 62d64a9..ba7e5c3 100644
--- 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
index 79432f9..ee0a8dd 100644
--- 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,