change strategy for testing CloseConnection races

Only pay for a task context switch when the code is under test.

Author: belm0

tests/test_connection.py

@@ -58,6 +58,7 @@
     open_websocket_url,
     serve_websocket,
     WebSocketServer,
+    WebSocketRequest,
     wrap_client_stream,
     wrap_server_stream
 )
@@ -906,8 +907,9 @@ async def handler(request):
 async def test_close_race(nursery, autojump_clock):
     """server attempts close just as client disconnects (issue #96)"""
 
-    async def handler(request):
+    async def handler(request: WebSocketRequest):
         ws = await request.accept()
+        ws._for_testing_peer_closed_connection = trio.Event()
         await ws.send_message('foo')
         await ws._for_testing_peer_closed_connection.wait()
         # with bug, this would raise ConnectionClosed from websocket internal task

trio_websocket/_impl.py

@@ -9,7 +9,7 @@
 import ssl
 import struct
 import urllib.parse
-from typing import List
+from typing import List, Optional
 
 import trio
 import trio.abc
@@ -744,9 +744,10 @@ def __init__(self, stream, ws_connection, *, host=None, path=None,
         # Set once a WebSocket closed handshake takes place, i.e after a close
         # frame has been sent and a close frame has been received.
         self._close_handshake = trio.Event()
-        # Set immediately upon receiving closed event from peer.  Used to
-        # test close race conditions between client and server.
-        self._for_testing_peer_closed_connection = trio.Event()
+        # When not None, set immediately upon receiving CloseConnection from peer,
+        # followed by sleep(0) for a context switch.  Used to test close race
+        # conditions between client and server.
+        self._for_testing_peer_closed_connection: Optional[trio.Event] = None
 
     @property
     def closed(self):
@@ -1095,8 +1096,9 @@ async def _handle_close_connection_event(self, event):
 
         :param wsproto.events.CloseConnection event:
         '''
-        self._for_testing_peer_closed_connection.set()
-        await trio.sleep(0)
+        if self._for_testing_peer_closed_connection:
+            self._for_testing_peer_closed_connection.set()
+            await trio.sleep(0)
         if self._wsproto.state == ConnectionState.REMOTE_CLOSING:
             await self._send(event.response())
         await self._close_web_socket(event.code, event.reason or None)