diff --git a/tests/unit_test/apis/impl/controller_test.py b/tests/unit_test/apis/impl/controller_test.py index 07af8846b2..e39accd28c 100644 --- 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