Fix controller unit test timing (#1936)

NVIDIA · Aug 21, 2023 · 95bf4af · 95bf4af
1 parent 6750969
commit 95bf4af
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/tests/unit_test/apis/impl/controller_test.py b/tests/unit_test/apis/impl/controller_test.py
@@ -691,6 +691,7 @@ def test_task_done_cb(self, method, num_clients, task_name, input_data, cb, expe
                     )
         if task_complete == "timeout":
             time.sleep(timeout)
+            controller._check_tasks()
             assert task.completion_status == TaskCompletionStatus.TIMEOUT
         elif task_complete == "cancel":
             controller.cancel_task(task)
@@ -1257,7 +1258,7 @@ def test_task_exit_quickly_when_all_responses_received(self, method, min_respons
                 client=client, task_name="__test_task", task_id=client_task_id, result=result, fl_ctx=fl_ctx
             )
 
-        time.sleep(0.5)
+        controller._check_tasks()
         assert controller.get_num_standing_tasks() == 0
         assert task.completion_status == TaskCompletionStatus.OK
         launch_thread.join()
@@ -1317,22 +1318,22 @@ def _process_task_request_test_cases():
     client_names = [c.name for c in clients]
 
     dynamic_targets_cases = [
-        (clients[1:], clients[0], True, 1, 0, False, [clients[1].name, clients[2].name, clients[0].name]),
-        (clients[1:], clients[1], True, 1, 0, True, client_names[1:]),
-        (clients[1:], clients[2], True, 1, 0, False, client_names[1:]),
-        ([clients[0]], clients[1], True, 1, 0, False, [clients[0].name, clients[1].name]),
+        (clients[1:], clients[0], True, 2, 0, False, [clients[1].name, clients[2].name, clients[0].name]),
+        (clients[1:], clients[1], True, 2, 0, True, client_names[1:]),
+        (clients[1:], clients[2], True, 2, 0, False, client_names[1:]),
+        ([clients[0]], clients[1], True, 2, 0, False, [clients[0].name, clients[1].name]),
         ([clients[0]], clients[1], True, 1, 2, False, [clients[0].name]),
         ([clients[0], clients[0]], clients[0], True, 1, 0, True, [clients[0].name, clients[0].name]),
         (None, clients[0], True, 1, 0, True, [clients[0].name]),
     ]
 
     static_targets_cases = [
-        (clients[1:], clients[0], False, 1, 0, False, client_names[1:]),
-        (clients[1:], clients[1], False, 1, 0, True, client_names[1:]),
-        (clients[1:], clients[2], False, 1, 0, False, client_names[1:]),
-        (clients[1:], clients[0], False, 1, 2, False, client_names[1:]),
-        (clients[1:], clients[1], False, 1, 2, True, client_names[1:]),
-        (clients[1:], clients[2], False, 1, 2, True, client_names[1:]),
+        (clients[1:], clients[0], False, 2, 0, False, client_names[1:]),
+        (clients[1:], clients[1], False, 2, 0, True, client_names[1:]),
+        (clients[1:], clients[2], False, 2, 0, False, client_names[1:]),
+        (clients[1:], clients[0], False, 1, 1.5, False, client_names[1:]),
+        (clients[1:], clients[1], False, 1, 1.5, True, client_names[1:]),
+        (clients[1:], clients[2], False, 1, 1.5, True, client_names[1:]),
     ]
 
     return dynamic_targets_cases + static_targets_cases