tasks: Add helper for identifying status as final

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

cogstack_model_gateway/common/tasks.py

@@ -18,6 +18,9 @@ class Status(str, Enum):
     FAILED = "failed"
     REQUEUED = "requeued"
 
+    def is_final(self) -> bool:
+        return self in {self.SUCCEEDED, self.FAILED}
+
 
 class Task(SQLModel, table=True):
     uuid: str = Field(default_factory=lambda: str(uuid.uuid4()), primary_key=True)

tests/unit/common/test_tasks.py

@@ -26,6 +26,13 @@ def task_manager(db_manager: Session) -> TaskManager:
     return TaskManager(db_manager)
 
 
+def test_status_enum() -> None:
+    for status in Status:
+        assert isinstance(status, Status)
+        assert isinstance(status.value, str)
+        assert status.is_final() == (status in {Status.SUCCEEDED, Status.FAILED})
+
+
 def test_create_task(task_manager: TaskManager) -> None:
     """Test creating a task with different statuses."""
     task_uuid = task_manager.create_task(status=Status.PENDING)