diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index b6426230..4e983bbe 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -574,7 +574,7 @@ async def __aenter__(self) -> ClientConnection: if hasattr(self, "connection"): raise RuntimeError("connect() isn't reentrant") self.connection = await self - return self.connection + return await self.connection.__aenter__() async def __aexit__( self, @@ -583,7 +583,7 @@ async def __aexit__( traceback: TracebackType | None, ) -> None: try: - await self.connection.close() + return await self.connection.__aexit__(exc_type, exc_value, traceback) finally: del self.connection diff --git a/src/websockets/asyncio/server.py b/src/websockets/asyncio/server.py index 7a168d04..58c1951d 100644 --- a/src/websockets/asyncio/server.py +++ b/src/websockets/asyncio/server.py @@ -744,12 +744,10 @@ async def protocol_handler(connection: ServerConnection) -> None: server.all_connections.add(connection) connection.start_keepalive() try: - await handler(connection) + async with connection: + await handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - await connection.close(CloseCode.INTERNAL_ERROR) - else: - await connection.close() finally: server.all_connections.discard(connection) diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 95fe53ab..393fc8e2 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -521,7 +521,7 @@ def __enter__(self) -> ClientConnection: raise RuntimeError("connect() isn't reentrant") self.connection = self.connect() self.connection.pending_legacy_warning = False - return self.connection + return self.connection.__enter__() def __exit__( self, @@ -530,7 +530,7 @@ def __exit__( exc_traceback: TracebackType | None, ) -> None: try: - self.connection.close() + return self.connection.__exit__(exc_type, exc_value, exc_traceback) finally: del self.connection diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index f0f3e058..d9ffac45 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -770,12 +770,10 @@ def protocol_select_subprotocol( server.all_connections.add(connection) connection.start_keepalive() try: - handler(connection) + with connection: + handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - connection.close(CloseCode.INTERNAL_ERROR) - else: - connection.close() finally: with server.lock: server.all_connections.discard(connection) diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 069fe4b2..a18cf7f2 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -549,7 +549,7 @@ async def __aenter__(self) -> ClientConnection: await self.__aenter_nursery__() try: self.connection = await self.connect(self.nursery) - return self.connection + return await self.connection.__aenter__() except BaseException as exc: await self.__aexit_nursery__(type(exc), exc, exc.__traceback__) raise AssertionError("expected __aexit_nursery__ to re-raise the exception") @@ -562,7 +562,7 @@ async def __aexit__( ) -> None: try: try: - await self.connection.aclose() + return await self.connection.__aexit__(exc_type, exc_value, traceback) finally: del self.connection finally: diff --git a/src/websockets/trio/server.py b/src/websockets/trio/server.py index 4e104c77..22531ed3 100644 --- a/src/websockets/trio/server.py +++ b/src/websockets/trio/server.py @@ -639,12 +639,10 @@ def protocol_select_subprotocol( server.all_connections.add(connection) connection.start_keepalive() try: - await handler(connection) + async with connection: + await handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - await connection.aclose(CloseCode.INTERNAL_ERROR) - else: - await connection.aclose() finally: server.all_connections.discard(connection) diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index f8f56308..46bbb59b 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -60,6 +60,21 @@ async def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + async def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + async with serve(*args) as server: + async with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + async def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + async with serve(*args) as server: + with self.assertRaises(RuntimeError): + async with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + async def test_direct_connection(self): """Client connects to server directly.""" async with serve(*args) as server: diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index e812fcc4..93c81d74 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -67,6 +67,21 @@ def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + with run_server() as server: + with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + with run_server() as server: + with self.assertRaises(RuntimeError): + with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + def test_direct_connection(self): """Client connects to server directly.""" with run_server() as server: diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 8d9d3774..b318a83a 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -62,6 +62,21 @@ async def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + async def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + async with run_server() as server: + async with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + async def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + async with run_server() as server: + with self.assertRaises(RuntimeError): + async with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" async with run_server() as server: