From cf9d22f3bc56789ee95621a0231e4092d04f3a24 Mon Sep 17 00:00:00 2001
From: Mikhail Sveshnikov <mike0sv@gmail.com>
Date: Wed, 15 Nov 2023 19:23:38 +0300
Subject: [PATCH] Counter panels for TestSuites (#843)

---
 src/evidently/pydantic_utils.py | 10 ++++++-
 src/evidently/ui/dashboards.py  | 48 +++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/evidently/pydantic_utils.py b/src/evidently/pydantic_utils.py
index cfffe0a637..2a8ac8d018 100644
--- a/src/evidently/pydantic_utils.py
+++ b/src/evidently/pydantic_utils.py
@@ -138,9 +138,17 @@ def __evidently_dependencies__(self):
 
 
 class EnumValueMixin(BaseModel):
+    def _to_enum_value(self, key, value):
+        field = self.__fields__[key]
+        if not issubclass(field.type_, Enum):
+            return value
+        if isinstance(value, list):
+            return [v.value if isinstance(v, Enum) else v for v in value]
+        return value.value if isinstance(value, Enum) else value
+
     def dict(self, *args, **kwargs) -> "DictStrAny":
         res = super().dict(*args, **kwargs)
-        return {k: v.value if isinstance(v, Enum) else v for k, v in res.items()}
+        return {k: self._to_enum_value(k, v) for k, v in res.items()}
 
 
 class ExcludeNoneMixin(BaseModel):
diff --git a/src/evidently/ui/dashboards.py b/src/evidently/ui/dashboards.py
index 4898d6a92b..189d8fc1d3 100644
--- a/src/evidently/ui/dashboards.py
+++ b/src/evidently/ui/dashboards.py
@@ -1,6 +1,7 @@
 import abc
 import datetime
 import traceback
+import typing
 import uuid
 from collections import Counter
 from collections import defaultdict
@@ -424,6 +425,53 @@ def _to_period(self, timestamp: datetime.datetime) -> datetime.datetime:
         return pd.Series([timestamp], name="dt").dt.to_period(self.time_agg)[0]
 
 
+class DashboardPanelTestSuiteCounter(DashboardPanel):
+    agg: CounterAgg = CounterAgg.NONE
+    filter: ReportFilter = ReportFilter(metadata_values={}, tag_values=[], include_test_suites=True)
+    test_filters: List[TestFilter] = []
+    statuses: List[TestStatus] = [TestStatus.SUCCESS]
+
+    def _iter_statuses(self, reports: Iterable[ReportBase]):
+        for report in reports:
+            if not self.filter.filter(report):
+                continue
+            if not isinstance(report, TestSuite):
+                continue
+            if self.test_filters:
+                for test_filter in self.test_filters:
+                    yield report.timestamp, test_filter.get(report).values()
+            else:
+                yield report.timestamp, TestFilter().get(report).values()
+
+    def build_widget(self, reports: Iterable[ReportBase]) -> BaseWidgetInfo:
+        if self.agg == CounterAgg.NONE:
+            statuses, postfix = self._build_none(reports)
+        elif self.agg == CounterAgg.LAST:
+            statuses, postfix = self._build_last(reports)
+        else:
+            raise ValueError(f"TestSuite Counter does not support agg {self.agg}")
+
+        total = sum(statuses.values())
+        value = sum(statuses[s] for s in self.statuses)
+        statuses_join = ", ".join(s.value for s in self.statuses)
+        return counter(counters=[CounterData(f"{value}/{total} {statuses_join}{postfix}", self.title)], size=self.size)
+
+    def _build_none(self, reports: Iterable[ReportBase]) -> Tuple[Counter, str]:
+        statuses: typing.Counter[TestStatus] = Counter()
+        for _, values in self._iter_statuses(reports):
+            statuses.update(values)
+        return statuses, ""
+
+    def _build_last(self, reports: Iterable[ReportBase]) -> Tuple[Counter, str]:
+        last_ts = None
+        statuses: typing.Counter[TestStatus] = Counter()
+        for ts, values in self._iter_statuses(reports):
+            if last_ts is None or ts > last_ts:
+                last_ts = ts
+                statuses = Counter(values)
+        return statuses, f" ({last_ts})"
+
+
 class DashboardConfig(BaseModel):
     name: str
     panels: List[DashboardPanel]