Merge pull request #376 from nttcom/bugfix/ssvc-order

fix comparison SSVCPriority
nttcom · Oct 1, 2024 · 5ce3875 · 5ce3875
2 parents 2377874 + 9cda0b2
commit 5ce3875
Show file tree

Hide file tree

Showing 2 changed files with 139 additions and 4 deletions.
diff --git a/api/app/models.py b/api/app/models.py
@@ -1,7 +1,6 @@
 import enum
 import uuid
 from datetime import datetime
-from functools import total_ordering
 from typing import cast
 
 from sqlalchemy import ARRAY, JSON, ForeignKey, LargeBinary, String, Text, UniqueConstraint
@@ -203,17 +202,41 @@ class HumanImpactEnum(str, enum.Enum):
     VERY_HIGH = "very_high"
 
 
-@total_ordering
 class SSVCDeployerPriorityEnum(str, enum.Enum):
     # https://certcc.github.io/SSVC/howto/deployer_tree/#deployer-decision-outcomes
     IMMEDIATE = "immediate"
     OUT_OF_CYCLE = "out_of_cycle"
     SCHEDULED = "scheduled"
     DEFER = "defer"
 
+    @property
+    def orders_map(self):
+        return {
+            SSVCDeployerPriorityEnum.IMMEDIATE: 1,
+            SSVCDeployerPriorityEnum.OUT_OF_CYCLE: 2,
+            SSVCDeployerPriorityEnum.SCHEDULED: 3,
+            SSVCDeployerPriorityEnum.DEFER: 4,
+        }
+
     def __lt__(self, other):
-        orders_map = {"immediate": 1, "out_of_cycle": 2, "scheduled": 3, "defer": 4}
-        return orders_map[self] < orders_map[other]
+        if not isinstance(other, SSVCDeployerPriorityEnum):
+            return NotImplemented
+        return self.orders_map[self] < self.orders_map[other]
+
+    def __le__(self, other):
+        if not isinstance(other, SSVCDeployerPriorityEnum):
+            return NotImplemented
+        return self.orders_map[self] <= self.orders_map[other]
+
+    def __gt__(self, other):
+        if not isinstance(other, SSVCDeployerPriorityEnum):
+            return NotImplemented
+        return self.orders_map[self] > self.orders_map[other]
+
+    def __ge__(self, other):
+        if not isinstance(other, SSVCDeployerPriorityEnum):
+            return NotImplemented
+        return self.orders_map[self] >= self.orders_map[other]
 
 
 # Base class

diff --git a/api/app/tests/small/test_ssvc_order.py b/api/app/tests/small/test_ssvc_order.py
@@ -0,0 +1,112 @@
+import pytest
+
+from app import models
+
+
+def test_ssvc_priority_enum_comparison():
+    assert models.SSVCDeployerPriorityEnum.DEFER == models.SSVCDeployerPriorityEnum.DEFER
+    assert models.SSVCDeployerPriorityEnum.DEFER != models.SSVCDeployerPriorityEnum.SCHEDULED
+    assert models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE < models.SSVCDeployerPriorityEnum.SCHEDULED
+    assert models.SSVCDeployerPriorityEnum.SCHEDULED <= models.SSVCDeployerPriorityEnum.DEFER
+    assert models.SSVCDeployerPriorityEnum.DEFER > models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE
+    assert models.SSVCDeployerPriorityEnum.DEFER >= models.SSVCDeployerPriorityEnum.IMMEDIATE
+
+
+def test_ssvc_priority_enum_comparison_with_different_type():
+    with pytest.raises(TypeError):
+        assert models.SSVCDeployerPriorityEnum.DEFER < 3
+    with pytest.raises(TypeError):
+        assert models.SSVCDeployerPriorityEnum.DEFER > 3
+    with pytest.raises(TypeError):
+        assert models.SSVCDeployerPriorityEnum.DEFER <= 3
+    with pytest.raises(TypeError):
+        assert models.SSVCDeployerPriorityEnum.DEFER <= 3
+
+
+@pytest.mark.parametrize(
+    "ssvc_priorityA, ssvc_priorityB, expected",
+    [
+        (
+            models.SSVCDeployerPriorityEnum.DEFER,
+            models.SSVCDeployerPriorityEnum.DEFER,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.DEFER,
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.DEFER,
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.DEFER,
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            models.SSVCDeployerPriorityEnum.DEFER,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            models.SSVCDeployerPriorityEnum.DEFER,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            False,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            models.SSVCDeployerPriorityEnum.DEFER,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            models.SSVCDeployerPriorityEnum.SCHEDULED,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            models.SSVCDeployerPriorityEnum.OUT_OF_CYCLE,
+            True,
+        ),
+        (
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            models.SSVCDeployerPriorityEnum.IMMEDIATE,
+            True,
+        ),
+    ],
+)
+def test_ssvc_priority_enum_le(ssvc_priorityA, ssvc_priorityB, expected):
+    assert (ssvc_priorityA <= ssvc_priorityB) == expected