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