From 3241edb1d251ea29079380c86997660c43f5e35d Mon Sep 17 00:00:00 2001 From: Anda Zhou <83614683+azhou-determined@users.noreply.github.com> Date: Fri, 20 Sep 2024 08:18:27 -0700 Subject: [PATCH] fix: fix flaky generic task pause test (#9962) --- e2e_tests/tests/task/test_generic_tasks.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/e2e_tests/tests/task/test_generic_tasks.py b/e2e_tests/tests/task/test_generic_tasks.py index 3a39a65e817..907a3afbed3 100644 --- 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)