Cleanup: another attempt at fixing cancellation tests

covalent/executor/schemas.py

@@ -22,28 +22,27 @@
 Types defining the runner-executor interface
 """
 
-from enum import Enum
 from typing import Dict, List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
 
 from covalent._shared_files.schemas.asset import AssetUpdate
-from covalent._shared_files.util_classes import RESULT_STATUS
-
-
-# Valid terminal statuses
-class TerminalStatus(str, Enum):
-    CANCELLED = RESULT_STATUS.CANCELLED
-    COMPLETED = RESULT_STATUS.COMPLETED
-    FAILED = RESULT_STATUS.FAILED
+from covalent._shared_files.util_classes import RESULT_STATUS, Status
 
 
 class TaskUpdate(BaseModel):
     dispatch_id: str
     node_id: int
-    status: TerminalStatus
+    status: Status
     assets: Dict[str, AssetUpdate]
 
+    @validator("status")
+    def validate_status(cls, v):
+        if RESULT_STATUS.is_terminal(v):
+            return v
+        else:
+            raise ValueError(f"Illegal status update {v}")
+
 
 class TaskSpec(BaseModel):
     function_id: int

covalent_dispatcher/_core/runner_ng.py

@@ -228,7 +228,7 @@ async def _get_task_result(task_group_metadata: Dict, data: Any):
         node_results = []
         for task_result in task_group_results:
             task_id = task_result.node_id
-            status = task_result.status.value
+            status = task_result.status
             await am.download_assets_for_node(dispatch_id, task_id, task_result.assets)
 
             node_result = datamgr.generate_node_result(