From 7966e2ccb533f4876807c053d64a892c1b2b7782 Mon Sep 17 00:00:00 2001
From: Drew Ripberger <drew.ripberger@gmail.com>
Date: Tue, 30 Jan 2024 13:59:53 -0500
Subject: [PATCH] use RegressionResult model in indicator

---
 matrix_benchmarking/models.py              |  7 +++--
 matrix_benchmarking/regression/__init__.py | 34 +++++++++++++---------
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/matrix_benchmarking/models.py b/matrix_benchmarking/models.py
index 258b059d..0357bfc9 100644
--- a/matrix_benchmarking/models.py
+++ b/matrix_benchmarking/models.py
@@ -133,8 +133,9 @@ def tostr(self):
 
 class RegressionResult(ExclusiveModel):
     kpi: str
+    setting: str
     indicator: str
     status: int
-    direction: Optional[int]
-    explanation: Optional[str]
-    details: Optional[dict[str, str]]
+    direction: Optional[int] = Field(default=None)
+    explanation: Optional[str] = Field(default=None)
+    details: Optional[dict[str, str]] = Field(default=None)
diff --git a/matrix_benchmarking/regression/__init__.py b/matrix_benchmarking/regression/__init__.py
index b92c8bed..951836ce 100644
--- a/matrix_benchmarking/regression/__init__.py
+++ b/matrix_benchmarking/regression/__init__.py
@@ -8,6 +8,7 @@
 from typing import Optional, Callable
 
 import matrix_benchmarking.common as common
+import matrix_benchmarking.models as models
 
 def get_from_path(d, path):
     return reduce(dict.get, path.split("."), d)
@@ -29,7 +30,6 @@ def __init__(
         self.explanation = explanation
         self.details = details
 
-
 class RegressionIndicator:
     """
     Assume the matrix that is passed in contains a prefiltered combination of settings,
@@ -82,13 +82,17 @@ def __init__(
     def get_name(self):
         return "UndefinedRegressionIndicator"
 
-    def analyze(self) -> list[dict]:
-
-        if not self.new_payload:
-            return [{"result": None, "kpi": None, "regression": vars(RegressionStatus(0, explanation="Not enough new data"))}]
+    def analyze(self) -> list[models.RegressionResult]:
 
-        if not self.lts_payloads:
-            return [{"result": None, "kpi": None, "regression": vars(RegressionStatus(0, explanation="Not enough LTS data"))}]
+        if not self.new_payload or not self.lts_payloads:
+            return [
+                models.RegressionResult(
+                    kpi="",
+                    setting="" if not self.x_var else self.x_var,
+                    indicator=self.get_name(),
+                    status=0
+                )
+            ]
 
         regression_results = []
 
@@ -107,12 +111,16 @@ def analyze(self) -> list[dict]:
                     continue
 
 
-            raw_results = self.regression_test(curr_values, lts_values)
-            stats = {
-                "kpi": kpi,
-                "indicator": self.get_name()
-            }
-            regression_results.append({**stats, **raw_results})
+            raw_results: RegressionStatus = self.regression_test(curr_values, lts_values)
+            result = models.RegressionResult(
+                kpi=kpi,
+                setting="" if not self.x_var else self.x_var,
+                indicator=self.get_name(),
+                direction=raw_results.direction,
+                explanation=raw_results.explanation,
+                details=raw_results.details
+            )
+            regression_results.append(result)
 
         return regression_results