fix: fix flaky generic task pause test (#9962)

determined-ai · Sep 20, 2024 · 3241edb · 3241edb
1 parent 43556e9
commit 3241edb
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/e2e_tests/tests/task/test_generic_tasks.py b/e2e_tests/tests/task/test_generic_tasks.py
@@ -231,15 +231,12 @@ def test_pause_and_unpause_generic_task() -> None:
 
     detproc.check_call(sess, command)
 
-    pause_resp = bindings.get_GetTask(sess, taskId=task_resp.taskId)
-    assert pause_resp.task.taskState == bindings.v1GenericTaskState.PAUSED
+    task.wait_for_task_state(sess, task_resp.taskId, bindings.v1GenericTaskState.PAUSED)
 
     # Unpause task
     command = ["det", "-m", conf.make_master_url(), "task", "unpause", task_resp.taskId]
 
     detproc.check_call(sess, command)
 
-    unpause_resp = bindings.get_GetTask(sess, taskId=task_resp.taskId)
-    assert unpause_resp.task.taskState == bindings.v1GenericTaskState.ACTIVE
-
+    task.wait_for_task_state(sess, task_resp.taskId, bindings.v1GenericTaskState.ACTIVE)
     task.wait_for_task_state(sess, task_resp.taskId, bindings.v1GenericTaskState.COMPLETED)