fix sender shutdown

wandb · Jan 11, 2021 · 85dc8c8 · 85dc8c8
1 parent fe3a5dc
commit 85dc8c8
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/tests/test_sender.py b/tests/test_sender.py
@@ -178,7 +178,15 @@ def stop_backend(
     mocked_run, hm, sm, sender, start_handle_thread, start_send_thread,
 ):
     def stop_backend_func():
-        sender.communicate_exit(0, timeout=5)
+        sender.publish_exit(0)
+        for _ in range(10):
+            ret = sender.communicate_poll_exit()
+            assert ret, "poll exit timedout"
+            done = ret.response.poll_exit_response.done
+            if done:
+                break
+            time.sleep(1)
+        assert done, "backend didnt shutdown"
 
     yield stop_backend_func