refactor MockBug in preperation for adding `status` field to `Node`

[utils.git] / src / budget_sync / util.py

diff --git a/src/budget_sync/util.py b/src/budget_sync/util.py
index ed1c300d38df42f01e836dc10dbec826b8366886..20e7da690b5ab2ea10ffee7c895929c14365b61b 100644 (file)
--- a/src/budget_sync/util.py
+++ b/src/budget_sync/util.py
@@ -1,6 +1,34 @@
 from bugzilla import Bugzilla
 from bugzilla.bug import Bug
-from typing import Iterator
+from typing import Iterator, Union
+from enum import Enum
+
+
+class BugStatus(Enum):
+    UNCONFIRMED = "UNCONFIRMED"
+    CONFIRMED = "CONFIRMED"
+    IN_PROGRESS = "IN_PROGRESS"
+    DEFERRED = "DEFERRED"
+    RESOLVED = "RESOLVED"
+    VERIFIED = "VERIFIED"
+    PAYMENTPENDING = "PAYMENTPENDING"
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return f"BugStatus.{self.value}"
+
+    @staticmethod
+    def cast(v: Union[str, "BugStatus"],
+             unknown_allowed: bool = False) -> Union[str, "BugStatus"]:
+        s = str(v)
+        try:
+            return BugStatus(s)
+        except ValueError:
+            if unknown_allowed:
+                return s
+            raise
 
 
 def all_bugs(bz: Bugzilla) -> Iterator[Bug]: