diff --git a/kmip/core/enums.py b/kmip/core/enums.py index 0f155c1..fb3fdce 100644 --- a/kmip/core/enums.py +++ b/kmip/core/enums.py @@ -20,6 +20,34 @@ import copy import enum +class OrderedEnum(enum.Enum): + """ + An ordered variant of the Enum class that allows for comparisons. + + Taken from: https://docs.python.org/3/library/enum.html#orderedenum + """ + + def __ge__(self, other): + if self.__class__ is other.__class__: + return self.value >= other.value + return NotImplemented + + def __gt__(self, other): + if self.__class__ is other.__class__: + return self.value > other.value + return NotImplemented + + def __le__(self, other): + if self.__class__ is other.__class__: + return self.value <= other.value + return NotImplemented + + def __lt__(self, other): + if self.__class__ is other.__class__: + return self.value < other.value + return NotImplemented + + class AdjustmentType(enum.Enum): # KMIP 2.0 INCREMENT = 0x00000001 @@ -508,7 +536,7 @@ class KeyWrapType(enum.Enum): AS_REGISTERED = 0x00000002 -class KMIPVersion(enum.Enum): +class KMIPVersion(OrderedEnum): KMIP_1_0 = 1.0 KMIP_1_1 = 1.1 KMIP_1_2 = 1.2 diff --git a/kmip/tests/unit/core/test_enums.py b/kmip/tests/unit/core/test_enums.py index 6d9b3e6..517af20 100644 --- a/kmip/tests/unit/core/test_enums.py +++ b/kmip/tests/unit/core/test_enums.py @@ -18,6 +18,67 @@ import testtools from kmip.core import enums +class TestOrderedEnum(testtools.TestCase): + + def setUp(self): + super(TestOrderedEnum, self).setUp() + + def tearDown(self): + super(TestOrderedEnum, self).tearDown() + + def test_greater_than_or_equal(self): + self.assertTrue( + enums.KMIPVersion.KMIP_2_0 >= enums.KMIPVersion.KMIP_1_0 + ) + self.assertFalse( + enums.KMIPVersion.KMIP_1_0 >= enums.KMIPVersion.KMIP_2_0 + ) + + self.assertEquals( + NotImplemented, + enums.KMIPVersion.KMIP_2_0.__ge__(enums.WrappingMethod.ENCRYPT) + ) + + def test_greater_than(self): + self.assertTrue( + enums.KMIPVersion.KMIP_1_3 > enums.KMIPVersion.KMIP_1_1 + ) + self.assertFalse( + enums.KMIPVersion.KMIP_1_1 > enums.KMIPVersion.KMIP_1_3 + ) + + self.assertEquals( + NotImplemented, + enums.KMIPVersion.KMIP_2_0.__gt__(enums.WrappingMethod.ENCRYPT) + ) + + def test_less_than_or_equal(self): + self.assertTrue( + enums.KMIPVersion.KMIP_1_3 <= enums.KMIPVersion.KMIP_1_4 + ) + self.assertFalse( + enums.KMIPVersion.KMIP_1_4 <= enums.KMIPVersion.KMIP_1_3 + ) + + self.assertEquals( + NotImplemented, + enums.KMIPVersion.KMIP_2_0.__le__(enums.WrappingMethod.ENCRYPT) + ) + + def test_less_than(self): + self.assertTrue( + enums.KMIPVersion.KMIP_1_3 < enums.KMIPVersion.KMIP_2_0 + ) + self.assertFalse( + enums.KMIPVersion.KMIP_2_0 < enums.KMIPVersion.KMIP_1_3 + ) + + self.assertEquals( + NotImplemented, + enums.KMIPVersion.KMIP_2_0.__lt__(enums.WrappingMethod.ENCRYPT) + ) + + class TestEnumUtilityFunctions(testtools.TestCase): def setUp(self):