From 3d7c9a1b0056cddfa8a750c049c48238f560ab69 Mon Sep 17 00:00:00 2001
From: moshemorad <moshe@trustmi.ai>
Date: Mon, 11 Nov 2024 13:07:38 +0200
Subject: [PATCH] Handle resources allocation with scientific notation better
 (#361)

---
 .gitignore                                |  1 +
 robusta_krr/utils/resource_units.py       |  9 +++---
 tests/models/test_resource_allocations.py | 36 +++++++++++++++++++++++
 3 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 tests/models/test_resource_allocations.py

diff --git a/.gitignore b/.gitignore
index 02d0ed42..4a103fe0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -133,3 +133,4 @@ dmypy.json
 .DS_Store
 robusta_lib
 .idea
+.vscode
\ No newline at end of file
diff --git a/robusta_krr/utils/resource_units.py b/robusta_krr/utils/resource_units.py
index 9fd290dc..dbda08bb 100644
--- a/robusta_krr/utils/resource_units.py
+++ b/robusta_krr/utils/resource_units.py
@@ -1,6 +1,6 @@
 from typing import Literal, Union
 
-UNITS = {
+UNITS: dict[str, float] = {
     "m": 0.001,
     "Ki": 1024,
     "Mi": 1024**2,
@@ -23,15 +23,14 @@ def parse(x: str, /) -> Union[float, int]:
     for unit, multiplier in UNITS.items():
         if x.endswith(unit):
             return float(x[: -len(unit)]) * multiplier
-    if "." in x:
-        return float(x)
-    return int(x)
+
+    return float(x)
 
 
 def get_base(x: str, /) -> Literal[1024, 1000]:
     """Returns the base of the unit."""
 
-    for unit, multiplier in UNITS.items():
+    for unit, _ in UNITS.items():
         if x.endswith(unit):
             return 1024 if unit in ["Ki", "Mi", "Gi", "Ti", "Pi", "Ei"] else 1000
     return 1000 if "." in x else 1024
diff --git a/tests/models/test_resource_allocations.py b/tests/models/test_resource_allocations.py
new file mode 100644
index 00000000..49ca949f
--- /dev/null
+++ b/tests/models/test_resource_allocations.py
@@ -0,0 +1,36 @@
+from typing import Union
+
+import pytest
+
+from robusta_krr.core.models.allocations import ResourceAllocations, ResourceType
+
+
+@pytest.mark.parametrize(
+    "cpu",
+    [
+        {"request": "5m", "limit": None},
+        {"request": 0.005, "limit": None},
+    ],
+)
+@pytest.mark.parametrize(
+    "memory",
+    [
+        {"request": 128974848, "limit": 128974848},
+        {"request": 128.974848e6, "limit": 128.974848e6},
+        {"request": "128.9748480M", "limit": "128.9748480M"},
+        {"request": "128974848000m", "limit": "128974848000m"},
+        {"request": "123Mi", "limit": "123Mi"},
+        {"request": "128974848e0", "limit": "128974848e0"},
+    ],
+)
+def test_resource_allocation_supported_formats(
+    cpu: dict[str, Union[str, int, float, None]], memory: dict[str, Union[str, int, float, None]]
+):
+    allocations = ResourceAllocations(
+        requests={ResourceType.CPU: cpu["request"], ResourceType.Memory: memory["request"]},
+        limits={ResourceType.CPU: cpu["limit"], ResourceType.Memory: memory["limit"]},
+    )
+    assert allocations.requests[ResourceType.CPU] == 0.005
+    assert allocations.limits[ResourceType.CPU] == None
+    assert (allocations.requests[ResourceType.Memory] // 1) == 128974848.0
+    assert (allocations.limits[ResourceType.Memory] // 1) == 128974848.0