Merge pull request #49 from robusta-dev/rework-severity-calculation

Rework severity calculation, fix score calculation

5 files changed, 199 insertions, 79 deletions

diff --git a/examples/custom_severity_calculator.py b/examples/custom_severity_calculator.py
new file mode 100644
index 0000000..5978bfc
--- /dev/null
+++ b/examples/custom_severity_calculator.py
@@ -0,0 +1,46 @@
+# This is an example on how to create your own custom formatter
+
+from __future__ import annotations
+
+from typing import Optional
+
+import robusta_krr
+from robusta_krr.api.models import Severity, ResourceType, register_severity_calculator
+
+
+@register_severity_calculator(ResourceType.CPU)
+def percentage_severity_calculator(
+    current: Optional[float], recommended: Optional[float], resource_type: ResourceType
+) -> Severity:
+    """
+    This is an example on how to create your own custom severity calculator
+    You can use this decorator to bind a severity calculator function to a resource type.
+    The function will be called with the current value, the recommended value and the resource type.
+    The function should return a Severity enum value.
+
+    If you have the same calculation for multiple resource types, you can use the `bind_calculator` decorator multiple times.
+    Then, the function will be called for each resource type and you can use the resource type to distinguish between them.
+
+    Keep in mind that you can not choose the strategy for the resource type using CLI - the last one created for the resource type will be used.
+    """
+
+    if current is None and recommended is None:
+        return Severity.GOOD
+    if current is None or recommended is None:
+        return Severity.WARNING
+
+    diff = abs(current - recommended) / current
+    if diff >= 0.5:
+        return Severity.CRITICAL
+    elif diff >= 0.25:
+        return Severity.WARNING
+    elif diff >= 0.1:
+        return Severity.OK
+    else:
+        return Severity.GOOD
+
+
+# Running this file will register the formatter and make it available to the CLI
+# Run it as `python ./custom_formatter.py simple --formater my_formatter`
+if __name__ == "__main__":
+    robusta_krr.run()
diff --git a/robusta_krr/api/models.py b/robusta_krr/api/models.py
index a0c4ea9..537168a 100644
--- a/robusta_krr/api/models.py
+++ b/robusta_krr/api/models.py
@@ -1,7 +1,8 @@
 from robusta_krr.core.abstract.strategies import HistoryData, ResourceHistoryData, ResourceRecommendation, RunResult
 from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType
 from robusta_krr.core.models.objects import K8sObjectData, PodData
-from robusta_krr.core.models.result import ResourceScan, Result, Severity
+from robusta_krr.core.models.result import ResourceScan, Result
+from robusta_krr.core.models.severity import Severity, register_severity_calculator
 
 __all__ = [
     "ResourceType",
@@ -11,6 +12,7 @@ __all__ = [
     "PodData",
     "Result",
     "Severity",
+    "register_severity_calculator",
     "ResourceScan",
     "ResourceRecommendation",
     "HistoryData",
diff --git a/robusta_krr/core/models/result.py b/robusta_krr/core/models/result.py
index 6e860b7..a72ff21 100644
--- a/robusta_krr/core/models/result.py
+++ b/robusta_krr/core/models/result.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import enum
-import itertools
 from datetime import datetime
 from typing import Any, Optional, Union
 
@@ -10,45 +8,7 @@ import pydantic as pd
 from robusta_krr.core.abstract import formatters
 from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType
 from robusta_krr.core.models.objects import K8sObjectData
-
-
-class Severity(str, enum.Enum):
-    """The severity of the scan."""
-
-    UNKNOWN = "UNKNOWN"
-    GOOD = "GOOD"
-    OK = "OK"
-    WARNING = "WARNING"
-    CRITICAL = "CRITICAL"
-
-    @property
-    def color(self) -> str:
-        return {
-            self.UNKNOWN: "dim",
-            self.GOOD: "green",
-            self.OK: "gray",
-            self.WARNING: "yellow",
-            self.CRITICAL: "red",
-        }[self]
-
-    @classmethod
-    def calculate(cls, current: RecommendationValue, recommended: RecommendationValue) -> Severity:
-        if isinstance(recommended, str) or isinstance(current, str):
-            return cls.UNKNOWN
-
-        if current is None and recommended is None:
-            return cls.OK
-        if current is None or recommended is None:
-            return cls.WARNING
-
-        diff = (current - recommended) / recommended
-
-        if diff > 1.0 or diff < -0.5:
-            return cls.CRITICAL
-        elif diff > 0.5 or diff < -0.25:
-            return cls.WARNING
-        else:
-            return cls.GOOD
+from robusta_krr.core.models.severity import Severity
 
 
 class Recommendation(pd.BaseModel):
@@ -88,7 +48,7 @@ class ResourceScan(pd.BaseModel):
                 current = getattr(object.allocations, selector).get(resource_type)
                 recommended = getattr(recommendation, selector).get(resource_type)
 
-                current_severity = Severity.calculate(current, recommended)
+                current_severity = Severity.calculate(current, recommended, resource_type)
 
                 getattr(recommendation_processed, selector)[resource_type] = Recommendation(
                     value=recommended, severity=current_severity
@@ -129,18 +89,8 @@ class Result(pd.BaseModel):
         return formatter(self)
 
     @staticmethod
-    def __percentage_difference(current: RecommendationValue, recommended: RecommendationValue) -> float:
-        """Get the percentage difference between two numbers.
-
-        Args:
-            current: The current value.
-            recommended: The recommended value.
-
-        Returns:
-            The percentage difference.
-        """
-
-        return 1
+    def __scan_cost(scan: ResourceScan) -> float:
+        return 0.7 if scan.severity == Severity.WARNING else 1 if scan.severity == Severity.CRITICAL else 0
 
     def __calculate_score(self) -> int:
         """Get the score of the result.
@@ -149,18 +99,19 @@ class Result(pd.BaseModel):
             The score of the result.
         """
 
-        total_diff = 0.0
-        for scan, resource_type in itertools.product(self.scans, ResourceType):
-            total_diff += self.__percentage_difference(
-                scan.object.allocations.requests[resource_type], scan.recommended.requests[resource_type]
-            )
-            total_diff += self.__percentage_difference(
-                scan.object.allocations.limits[resource_type], scan.recommended.limits[resource_type]
-            )
-
-        if len(self.scans) == 0:
-            return 0
-
-        return int(
-            max(0, round(100 - total_diff / len(self.scans) / len(ResourceType) / 50, 2))
-        )  # 50 is just a constant
+        score = sum(self.__scan_cost(scan) for scan in self.scans)
+        return int((len(self.scans) - score) / len(self.scans) * 100)
+
+    @property
+    def score_letter(self) -> str:
+        return (
+            "F"
+            if self.score < 30
+            else "D"
+            if self.score < 55
+            else "C"
+            if self.score < 70
+            else "B"
+            if self.score < 90
+            else "A"
+        )
diff --git a/robusta_krr/core/models/severity.py b/robusta_krr/core/models/severity.py
new file mode 100644
index 0000000..ec1fad0
--- /dev/null
+++ b/robusta_krr/core/models/severity.py
@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import enum
+
+from typing import Callable, Optional
+from robusta_krr.core.models.allocations import RecommendationValue, ResourceType
+
+
+class Severity(str, enum.Enum):
+    """
+    The severity of the scan.
+
+    The severity is calculated based on the difference between the current value and the recommended value.
+    You can override the severity calculation function by using the `bind_calculator` decorator from the same module.
+    """
+
+    UNKNOWN = "UNKNOWN"
+    GOOD = "GOOD"
+    OK = "OK"
+    WARNING = "WARNING"
+    CRITICAL = "CRITICAL"
+
+    @property
+    def color(self) -> str:
+        return {
+            self.UNKNOWN: "dim",
+            self.GOOD: "green",
+            self.OK: "gray",
+            self.WARNING: "yellow",
+            self.CRITICAL: "red",
+        }[self]
+
+    @classmethod
+    def calculate(
+        cls, current: RecommendationValue, recommended: RecommendationValue, resource_type: ResourceType
+    ) -> Severity:
+        if isinstance(recommended, str) or isinstance(current, str):
+            return cls.UNKNOWN
+
+        return calculate_severity(current, recommended, resource_type)
+
+
+def register_severity_calculator(resource_type: ResourceType) -> Callable[[SeverityCalculator], SeverityCalculator]:
+    """
+    Bind a severity calculator function to a resource type.
+    Use this decorator to bind a severity calculator function to a resource type.
+
+    Example:
+    >>> @bind_severity_calculator(ResourceType.CPU)
+    >>> def cpu_severity_calculator(current: Optional[float], recommended: Optional[float], resource_type: ResourceType) -> Severity:
+    >>>     if current is None and recommended is None:
+    >>>         return Severity.GOOD
+    >>>     if current is None or recommended is None:
+    >>>         return Severity.WARNING
+    >>>
+    >>>     return Severity.CRITICAL if abs(current - recommended) >= 0.5 else Severity.GOOD
+    """
+
+    def decorator(func: SeverityCalculator) -> SeverityCalculator:
+        SEVERITY_CALCULATORS_REGISTRY[resource_type] = func
+        return func
+
+    return decorator
+
+
+SeverityCalculator = Callable[[Optional[float], Optional[float], ResourceType], Severity]
+SEVERITY_CALCULATORS_REGISTRY: dict[ResourceType, SeverityCalculator] = {}
+
+
+def calculate_severity(current: Optional[float], recommended: Optional[float], resource_type: ResourceType) -> Severity:
+    """
+    Calculate the severity of the scan based on the current value and the recommended value.
+
+    This function will use the severity calculator function that is bound to the resource type.
+    If there is no calculator function bound to the resource type, it will use the default severity calculator function.
+    """
+
+    return SEVERITY_CALCULATORS_REGISTRY.get(resource_type, default_severity_calculator)(
+        current, recommended, resource_type
+    )
+
+
+def default_severity_calculator(
+    current: Optional[float], recommended: Optional[float], resource_type: ResourceType
+) -> Severity:
+    return Severity.UNKNOWN
+
+
+@register_severity_calculator(ResourceType.CPU)
+def cpu_severity_calculator(
+    current: Optional[float], recommended: Optional[float], resource_type: ResourceType
+) -> Severity:
+    if current is None and recommended is None:
+        return Severity.GOOD
+    if current is None or recommended is None:
+        return Severity.WARNING
+
+    diff = abs(current - recommended)
+
+    if diff >= 0.5:
+        return Severity.CRITICAL
+    elif diff >= 0.25:
+        return Severity.WARNING
+    elif diff >= 0.1:
+        return Severity.OK
+    else:
+        return Severity.GOOD
+
+
+@register_severity_calculator(ResourceType.Memory)
+def memory_severity_calculator(
+    current: Optional[float], recommended: Optional[float], resource_type: ResourceType
+) -> Severity:
+    if current is None and recommended is None:
+        return Severity.GOOD
+    if current is None or recommended is None:
+        return Severity.WARNING
+
+    diff = abs(current - recommended) / 1024 / 1024
+
+    if diff >= 500:
+        return Severity.CRITICAL
+    elif diff >= 250:
+        return Severity.WARNING
+    elif diff >= 100:
+        return Severity.OK
+    else:
+        return Severity.GOOD
diff --git a/robusta_krr/formatters/table.py b/robusta_krr/formatters/table.py
index 9776d80..c510fc2 100644
--- a/robusta_krr/formatters/table.py
+++ b/robusta_krr/formatters/table.py
@@ -25,13 +25,7 @@ def _format_request_str(item: ResourceScan, resource: ResourceType, selector: st
     recommended = getattr(item.recommended, selector)[resource]
     severity = recommended.severity
 
-    return (
-        f"[{severity.color}]"
-        + _format(allocated)
-        + " -> "
-        + _format(recommended.value)
-        + f"[/{severity.color}]"
-    )
+    return f"[{severity.color}]" + _format(allocated) + " -> " + _format(recommended.value) + f"[/{severity.color}]"
 
 
 @formatters.register(rich_console=True)
@@ -50,8 +44,7 @@ def table(result: Result) -> Table:
         title=f"\n{result.description}\n" if result.description else None,
         title_justify="left",
         title_style="",
-        # TODO: Fix points calculation at [MAIN-270]
-        # caption=f"Scan result ({result.score} points)",
+        caption=f"{result.score} points - {result.score_letter}",
     )
 
     table.add_column("Number", justify="right", no_wrap=True)