diff --git a/kubernetes/e2e_test/test_utils.py b/kubernetes/e2e_test/test_utils.py index cc80e6f20..8405e7422 100644 --- a/kubernetes/e2e_test/test_utils.py +++ b/kubernetes/e2e_test/test_utils.py @@ -13,13 +13,15 @@ # under the License. import unittest +from decimal import Decimal from os import path import yaml -from kubernetes import utils, client +from kubernetes import client, utils from kubernetes.client.rest import ApiException from kubernetes.e2e_test import base +from kubernetes.utils import quantity class TestUtils(unittest.TestCase): @@ -563,3 +565,113 @@ def test_create_from_list_in_multi_resource_yaml_namespaced(self): name="mock-pod-1", namespace=self.test_namespace, body={}) app_api.delete_namespaced_deployment( name="mock", namespace=self.test_namespace, body={}) + + +class TestUtilsUnitTests(unittest.TestCase): + + def test_parse_quantity(self): + # == trivial returns == + self.assertEqual(quantity.parse_quantity(Decimal(1)), Decimal(1)) + self.assertEqual(quantity.parse_quantity(float(1)), Decimal(1)) + self.assertEqual(quantity.parse_quantity(1), Decimal(1)) + + # == exceptions == + self.assertRaises( + ValueError, lambda: quantity.parse_quantity("1000kb") + ) + self.assertRaises( + ValueError, lambda: quantity.parse_quantity("1000ki") + ) + self.assertRaises(ValueError, lambda: quantity.parse_quantity("1000foo")) + self.assertRaises(ValueError, lambda: quantity.parse_quantity("foo")) + + # == no suffix == + self.assertEqual(quantity.parse_quantity("1000"), Decimal(1000)) + + # == base 1024 == + self.assertEqual(quantity.parse_quantity("1Ki"), Decimal(1024)) + self.assertEqual(quantity.parse_quantity("1Mi"), Decimal(1024**2)) + self.assertEqual(quantity.parse_quantity("1Gi"), Decimal(1024**3)) + self.assertEqual(quantity.parse_quantity("1Ti"), Decimal(1024**4)) + self.assertEqual(quantity.parse_quantity("1Pi"), Decimal(1024**5)) + self.assertEqual(quantity.parse_quantity("1Ei"), Decimal(1024**6)) + self.assertEqual(quantity.parse_quantity("1024Ki"), Decimal(1024**2)) + self.assertEqual(quantity.parse_quantity("0.5Ki"), Decimal(512)) + + # == base 1000 == + self.assertAlmostEqual(quantity.parse_quantity("1n"), Decimal(0.000_000_001)) + self.assertAlmostEqual(quantity.parse_quantity("1u"), Decimal(0.000_001)) + self.assertAlmostEqual(quantity.parse_quantity("1m"), Decimal(0.001)) + self.assertEqual(quantity.parse_quantity("1k"), Decimal(1_000)) + self.assertEqual(quantity.parse_quantity("1M"), Decimal(1_000_000)) + self.assertEqual(quantity.parse_quantity("1G"), Decimal(1_000_000_000)) + self.assertEqual(quantity.parse_quantity("1T"), Decimal(1_000_000_000_000)) + self.assertEqual(quantity.parse_quantity("1P"), Decimal(1_000_000_000_000_000)) + self.assertEqual( + quantity.parse_quantity("1E"), Decimal(1_000_000_000_000_000_000)) + self.assertEqual(quantity.parse_quantity("1000k"), Decimal(1_000_000)) + self.assertEqual(quantity.parse_quantity("500k"), Decimal(500_000)) + + def test_format_quantity(self): + """Unit test for quantity.format_quantity. Testing the different SI suffixes and + function should return the expected string""" + + # == unknown suffixes == + self.assertRaises( + ValueError, lambda: quantity.format_quantity(Decimal(1_000), "kb") + ) + self.assertRaises( + ValueError, lambda: quantity.format_quantity(Decimal(1_000), "ki") + ) + self.assertRaises( + ValueError, lambda: quantity.format_quantity(Decimal(1_000), "foo") + ) + + # == no suffix == + self.assertEqual(quantity.format_quantity(Decimal(1_000), ""), "1000") + self.assertEqual(quantity.format_quantity(Decimal(1_000), None), "1000") + + # == base 1024 == + self.assertEqual(quantity.format_quantity(Decimal(1024), "Ki"), "1Ki") + self.assertEqual(quantity.format_quantity(Decimal(1024**2), "Mi"), "1Mi") + self.assertEqual(quantity.format_quantity(Decimal(1024**3), "Gi"), "1Gi") + self.assertEqual(quantity.format_quantity(Decimal(1024**4), "Ti"), "1Ti") + self.assertEqual(quantity.format_quantity(Decimal(1024**5), "Pi"), "1Pi") + self.assertEqual(quantity.format_quantity(Decimal(1024**6), "Ei"), "1Ei") + self.assertEqual(quantity.format_quantity(Decimal(1024**2), "Ki"), "1024Ki") + self.assertEqual(quantity.format_quantity(Decimal((1024**3) / 2), "Gi"), "0.5Gi") + # Decimal((1024**3)/3) are 0.3333333333333333148296162562Gi; expecting to + # be quantized to 0.3Gi + self.assertEqual( + quantity.format_quantity( + Decimal( + (1024**3) / 3), + "Gi", + quantize=Decimal(.5)), + "0.3Gi") + + # == base 1000 == + self.assertEqual(quantity.format_quantity(Decimal(0.000_000_001), "n"), "1n") + self.assertEqual(quantity.format_quantity(Decimal(0.000_001), "u"), "1u") + self.assertEqual(quantity.format_quantity(Decimal(0.001), "m"), "1m") + self.assertEqual(quantity.format_quantity(Decimal(1_000), "k"), "1k") + self.assertEqual(quantity.format_quantity(Decimal(1_000_000), "M"), "1M") + self.assertEqual(quantity.format_quantity(Decimal(1_000_000_000), "G"), "1G") + self.assertEqual( + quantity.format_quantity(Decimal(1_000_000_000_000), "T"), "1T" + ) + self.assertEqual( + quantity.format_quantity(Decimal(1_000_000_000_000_000), "P"), "1P" + ) + self.assertEqual( + quantity.format_quantity(Decimal(1_000_000_000_000_000_000), "E"), "1E" + ) + self.assertEqual(quantity.format_quantity(Decimal(1_000_000), "k"), "1000k") + # Decimal(1_000_000/3) are 333.3333333333333139307796955k; expecting to + # be quantized to 333k + self.assertEqual( + quantity.format_quantity( + Decimal(1_000_000 / 3), "k", quantize=Decimal(1000) + ), + "333k", + ) diff --git a/kubernetes/utils/quantity.py b/kubernetes/utils/quantity.py index 68e57d980..484f34719 100644 --- a/kubernetes/utils/quantity.py +++ b/kubernetes/utils/quantity.py @@ -13,6 +13,19 @@ # limitations under the License. from decimal import Decimal, InvalidOperation +_EXPONENTS = { + "n": -3, + "u": -2, + "m": -1, + "K": 1, + "k": 1, + "M": 2, + "G": 3, + "T": 4, + "P": 5, + "E": 6, +} + def parse_quantity(quantity): """ @@ -35,17 +48,14 @@ def parse_quantity(quantity): if isinstance(quantity, (int, float, Decimal)): return Decimal(quantity) - exponents = {"n": -3, "u": -2, "m": -1, "K": 1, "k": 1, "M": 2, - "G": 3, "T": 4, "P": 5, "E": 6} - quantity = str(quantity) number = quantity suffix = None if len(quantity) >= 2 and quantity[-1] == "i": - if quantity[-2] in exponents: + if quantity[-2] in _EXPONENTS: number = quantity[:-2] suffix = quantity[-2:] - elif len(quantity) >= 1 and quantity[-1] in exponents: + elif len(quantity) >= 1 and quantity[-1] in _EXPONENTS: number = quantity[:-1] suffix = quantity[-1:] @@ -68,8 +78,65 @@ def parse_quantity(quantity): if suffix == "ki": raise ValueError("{} has unknown suffix".format(quantity)) - if suffix[0] not in exponents: + if suffix[0] not in _EXPONENTS: raise ValueError("{} has unknown suffix".format(quantity)) - exponent = Decimal(exponents[suffix[0]]) + exponent = Decimal(_EXPONENTS[suffix[0]]) return number * (base ** exponent) + + +def format_quantity(quantity_value, suffix, quantize=None) -> str: + """ + Takes a decimal and produces a string value in kubernetes' canonical quantity form, + like "200Mi".Users can specify an additional decimal number to quantize the output. + + Example - Relatively increase pod memory limits: + + # retrieve my_pod + current_memory: Decimal = parse_quantity(my_pod.spec.containers[0].resources.limits.memory) + desired_memory = current_memory * 1.2 + desired_memory_str = format_quantity(desired_memory, suffix="Gi", quantize=Decimal(1)) + # patch pod with desired_memory_str + + 'quantize=Decimal(1)' ensures that the result does not contain any fractional digits. + + Supported SI suffixes: + base1024: Ki | Mi | Gi | Ti | Pi | Ei + base1000: n | u | m | "" | k | M | G | T | P | E + + See https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go + + Input: + quantity: Decimal. Quantity as a number which is supposed to converted to a string + with SI suffix. + suffix: string. The desired suffix/unit-of-measure of the output string + quantize: Decimal. Can be used to round/quantize the value before the string + is returned. Defaults to None. + + Returns: + string. Canonical Kubernetes quantity string containing the SI suffix. + + Raises: + ValueError if the SI suffix is not supported. + """ + + if not suffix: + return str(quantity_value) + + if suffix.endswith("i"): + base = 1024 + elif len(suffix) == 1: + base = 1000 + else: + raise ValueError(f"{quantity_value} has unknown suffix") + + if suffix == "ki": + raise ValueError(f"{quantity_value} has unknown suffix") + + if suffix[0] not in _EXPONENTS: + raise ValueError(f"{quantity_value} has unknown suffix") + + different_scale = quantity_value / Decimal(base ** _EXPONENTS[suffix[0]]) + if quantize: + different_scale = different_scale.quantize(quantize) + return str(different_scale) + suffix