Ensure listener closes in test_close_callback

rapidsai · Oct 11, 2023 · 4a4c315 · 4a4c315
1 parent 3bf5589
commit 4a4c315
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/test_endpoint.py b/tests/test_endpoint.py
@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 import ucp
@@ -16,6 +18,7 @@ async def server_node(ep):
             ep.set_close_callback(_close_callback)
         if server_close_callback is False:
             await ep.close()
+        listener.close()
 
     async def client_node(port):
         ep = await ucp.create_endpoint(
@@ -31,6 +34,8 @@ async def client_node(port):
         server_node,
     )
     await client_node(listener.port)
+    while not listener.closed():
+        await asyncio.sleep(0.01)
     assert closed[0] is True