From 32ae74fd455df39401fd44945d4d748b66f46c84 Mon Sep 17 00:00:00 2001
From: Moritz Kiemer <moritz.kiemer@checkmk.com>
Date: Fri, 20 Oct 2023 23:28:04 +0200
Subject: [PATCH] prediction: introduce Slice

Change-Id: Ie599cc6a132f63ad65e2c9087dd313aa433deadf
---
 cmk/utils/prediction/_grouping.py         | 49 +++++++++++------------
 cmk/utils/prediction/_plugin_interface.py | 12 +++---
 2 files changed, 28 insertions(+), 33 deletions(-)

diff --git a/cmk/utils/prediction/_grouping.py b/cmk/utils/prediction/_grouping.py
index 66f7c58267e..27a39d368c5 100644
--- a/cmk/utils/prediction/_grouping.py
+++ b/cmk/utils/prediction/_grouping.py
@@ -5,7 +5,8 @@
 
 import time
 from collections.abc import Callable, Mapping, Sequence
-from typing import Final, Literal, NamedTuple, NewType
+from dataclasses import dataclass
+from typing import Final, Literal, NamedTuple, NewType, Self
 
 Timegroup = NewType("Timegroup", str)
 
@@ -99,7 +100,6 @@ def time_slices(
     "Collect all slices back into the past until time horizon is reached"
     timestamp = int(timestamp)
     abs_begin = timestamp - horizon_seconds
-    slices = []
 
     # Note: due to the f**king DST, we can have several shifts between DST
     # and non-DST during a computation. Treatment is unfair on those longer
@@ -110,27 +110,24 @@ def time_slices(
     # midnight) and most likely before DST swap is applied.
 
     # Have fun understanding the tests for this function.
-    for begin in range(timestamp, abs_begin, -period_info.slice):
-        tg, start, end = get_timegroup_relative_time(begin, period_info)[:3]
-        if tg == timegroup:
-            slices.append((start, end))
-    return slices
-
-
-def get_timegroup_relative_time(
-    timestamp: int,
-    period_info: PeriodInfo,
-) -> tuple[Timegroup, int, int, int]:
-    """
-    Return:
-    timegroup: name of the group, like 'monday' or '12'
-    from_time: absolute epoch time of the first second of the
-    current slice.
-    until_time: absolute epoch time of the first second *not* in the slice
-    rel_time: seconds offset of now in the current slice
-    """
-    # Convert to local timezone
-    timegroup, rel_time = period_info.groupby(timestamp)
-    from_time = timestamp - rel_time
-    until_time = from_time + period_info.slice
-    return timegroup, from_time, until_time, rel_time
+    return [
+        slice_.interval
+        for begin in range(timestamp, abs_begin, -period_info.slice)
+        if (slice_ := Slice.from_timestamp(begin, period_info)).group == timegroup
+    ]
+
+
+@dataclass
+class Slice:
+    group: Timegroup
+    """Name of the group, like 'monday' or '12'"""
+    interval: tuple[int, int]
+    """Absolute epoch times of the first second of the slice
+    and the first second *not* in the slice."""
+
+    @classmethod
+    def from_timestamp(cls, timestamp: int, period_info: PeriodInfo) -> Self:
+        timegroup, rel_time = period_info.groupby(timestamp)
+        from_time = timestamp - rel_time
+        until_time = from_time + period_info.slice
+        return cls(group=timegroup, interval=(from_time, until_time))
diff --git a/cmk/utils/prediction/_plugin_interface.py b/cmk/utils/prediction/_plugin_interface.py
index 29cf1dbb872..7c6eecd00c3 100644
--- a/cmk/utils/prediction/_plugin_interface.py
+++ b/cmk/utils/prediction/_plugin_interface.py
@@ -13,7 +13,7 @@
 from cmk.utils.log import VERBOSE
 from cmk.utils.servicename import ServiceName
 
-from ._grouping import get_timegroup_relative_time, PREDICTION_PERIODS, Timegroup
+from ._grouping import PREDICTION_PERIODS, Slice, Timegroup
 from ._paths import PREDICTION_DIR
 from ._prediction import (
     compute_prediction,
@@ -73,9 +73,7 @@ def get_predictive_levels(
     now = int(time.time())
     period_info = PREDICTION_PERIODS[params.period]
 
-    timegroup, current_slice_start, current_slice_end, _rel_time = get_timegroup_relative_time(
-        now, period_info
-    )
+    current_slice = Slice.from_timestamp(now, period_info)
 
     prediction_store = PredictionStore(
         PREDICTION_DIR,
@@ -87,14 +85,14 @@ def get_predictive_levels(
     if (
         data_for_pred := _get_prediction(
             store=prediction_store,
-            timegroup=timegroup,
+            timegroup=current_slice.group,
             params=params,
         )
     ) is None:
         info = PredictionInfo(
-            name=timegroup,
+            name=current_slice.group,
             time=now,
-            range=(current_slice_start, current_slice_end),
+            range=current_slice.interval,
             dsname=dsname,
             params=params,
         )