Fix metric data computation

2 files changed, 17 insertions, 24 deletions

diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
index a409810..b0c7699 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
@@ -1,7 +1,7 @@
-import datetime
 from typing import Any, Optional
 
 from .base_metric import BaseMetricLoader
+from robusta_krr.core.abstract.strategies import Metric
 
 PrometheusSeries = Any
 
@@ -45,7 +45,9 @@ class BaseFilteredMetricLoader(BaseMetricLoader):
         # takes kubelet job if exists, return first job alphabetically if it doesn't
         for target_name in target_names:
             relevant_series = [
-                series for series in series_list_result if BaseFilteredMetricLoader.get_target_name(series) == target_name
+                series
+                for series in series_list_result
+                if BaseFilteredMetricLoader.get_target_name(series) == target_name
             ]
             relevant_kubelet_metric = [series for series in relevant_series if series["metric"].get("job") == "kubelet"]
             if len(relevant_kubelet_metric) == 1:
@@ -55,8 +57,6 @@ class BaseFilteredMetricLoader(BaseMetricLoader):
             return_list.append(sorted_relevant_series[0])
         return return_list
 
-    async def query_prometheus(
-        self, query: str, start_time: datetime.datetime, end_time: datetime.datetime, step: datetime.timedelta
-    ) -> list[PrometheusSeries]:
-        result = await super().query_prometheus(query, start_time, end_time, step)
+    async def query_prometheus(self, metric: Metric) -> list[PrometheusSeries]:
+        result = await super().query_prometheus(metric)
         return self.filter_prom_jobs_results(result)
diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
index e79e19a..de4cf59 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
@@ -30,15 +30,13 @@ class BaseMetricLoader(Configurable, abc.ABC):
     def _step_to_string(self, step: datetime.timedelta) -> str:
         return f"{int(step.total_seconds()) // 60}m"
 
-    async def query_prometheus(
-        self, query: str, start_time: datetime.datetime, end_time: datetime.datetime, step: datetime.timedelta
-    ) -> list[dict]:
+    async def query_prometheus(self, metric: Metric) -> list[dict]:
         return await asyncio.to_thread(
             self.prometheus.custom_query_range,
-            query=query,
-            start_time=start_time,
-            end_time=end_time,
-            step=self._step_to_string(step),
+            query=metric.query,
+            start_time=metric.start_time,
+            end_time=metric.end_time,
+            step=metric.step,
         )
 
     async def load_data(
@@ -46,27 +44,22 @@ class BaseMetricLoader(Configurable, abc.ABC):
     ) -> ResourceHistoryData:
         query = self.get_query(object)
         end_time = datetime.datetime.now()
-        result = await self.query_prometheus(
+        metric = Metric(
             query=query,
             start_time=end_time - period,
             end_time=end_time,
-            step=step,
+            step=self._step_to_string(step),
         )
+        result = await self.query_prometheus(metric)
 
         if result == []:
             self.warning(f"Prometheus returned no {self.__class__.__name__} metrics for {object}")
-            return ResourceHistoryData(query=query, data={})
+            return ResourceHistoryData(metric=metric, data={})
 
         return ResourceHistoryData(
-            metric=Metric(
-                query=query,
-                start_time=end_time - period,
-                end_time=end_time,
-                step=self._step_to_string(step),
-            ),
+            metric=metric,
             data={
-                pod_result['metric']['pod']: np.array(pod_result["values"], dtype=np.float64)
-                for pod_result in result
+                pod_result["metric"]["pod"]: np.array(pod_result["values"], dtype=np.float64) for pod_result in result
             },
         )