Don't finish test until handler exits

python-trio · Sep 22, 2020 · 150ddb4 · 150ddb4
1 parent 6ffd494
commit 150ddb4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/test_connection.py b/tests/test_connection.py
@@ -733,11 +733,14 @@ async def handler(request):
 
 
 async def test_server_sends_after_close(nursery):
+    done = trio.Event()
+
     async def handler(request):
         server_ws = await request.accept()
         with pytest.raises(ConnectionClosed):
             while True:
                 await server_ws.send_message('Hello from server')
+        done.set()
 
     server = await nursery.start(serve_websocket, handler, HOST, 0, None)
     stream = await trio.open_tcp_stream(HOST, server.port)
@@ -747,6 +750,7 @@ async def handler(request):
         for x in range(2):
             await client_ws.send_message('Hello from client')
         await stream.aclose()
+    await done.wait()
 
 
 async def test_server_does_not_close_handshake(nursery):