additional refactoring

opensource-observer · Dec 18, 2024 · 36d3480 · 36d3480
1 parent 5cb3c49
commit 36d3480
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 18 deletions.
diff --git a/warehouse/metrics_tools/compute/service.py b/warehouse/metrics_tools/compute/service.py
@@ -322,9 +322,7 @@ async def _notify_job_pending(self, job_id: str, input: JobSubmitRequest):
     async def _notify_job_running(self, job_id: str):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.JOB,
+            QueryJobUpdate.create_job_update(
                 payload=QueryJobStateUpdate(
                     status=QueryJobStatus.RUNNING,
                     has_remaining_tasks=True,
@@ -335,9 +333,7 @@ async def _notify_job_running(self, job_id: str):
     async def _notify_job_task_completed(self, job_id: str, task_id: str):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.TASK,
+            QueryJobUpdate.create_task_update(
                 payload=QueryJobTaskUpdate(
                     task_id=task_id,
                     status=QueryJobTaskStatus.SUCCEEDED,
@@ -350,9 +346,7 @@ async def _notify_job_task_failed(
     ):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.TASK,
+            QueryJobUpdate.create_task_update(
                 payload=QueryJobTaskUpdate(
                     task_id=task_id,
                     status=QueryJobTaskStatus.FAILED,
@@ -364,9 +358,7 @@ async def _notify_job_task_failed(
     async def _notify_job_task_cancelled(self, job_id: str, task_id: str):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.TASK,
+            QueryJobUpdate.create_task_update(
                 payload=QueryJobTaskUpdate(
                     task_id=task_id,
                     status=QueryJobTaskStatus.CANCELLED,
@@ -377,9 +369,7 @@ async def _notify_job_task_cancelled(self, job_id: str, task_id: str):
     async def _notify_job_completed(self, job_id: str):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.JOB,
+            QueryJobUpdate.create_job_update(
                 payload=QueryJobStateUpdate(
                     status=QueryJobStatus.COMPLETED,
                     has_remaining_tasks=False,
@@ -395,9 +385,7 @@ async def _notify_job_failed(
     ):
         await self._update_job_state(
             job_id,
-            QueryJobUpdate(
-                time=datetime.now(),
-                scope=QueryJobUpdateScope.JOB,
+            QueryJobUpdate.create_job_update(
                 payload=QueryJobStateUpdate(
                     status=QueryJobStatus.FAILED,
                     has_remaining_tasks=has_remaining_tasks,

diff --git a/warehouse/metrics_tools/compute/test_types.py b/warehouse/metrics_tools/compute/test_types.py
@@ -124,6 +124,26 @@
             False,
             1,
         ),
+        (
+            "should fail if a task is cancelled",
+            [
+                QueryJobUpdate.create_job_update(
+                    QueryJobStateUpdate(
+                        status=QueryJobStatus.RUNNING,
+                        has_remaining_tasks=True,
+                    ),
+                ),
+                QueryJobUpdate.create_task_update(
+                    QueryJobTaskUpdate(
+                        status=QueryJobTaskStatus.CANCELLED,
+                        task_id="task_id",
+                    )
+                ),
+            ],
+            QueryJobStatus.FAILED,
+            True,
+            0,
+        ),
     ],
 )
 def test_query_job_state(