From 1eb023c85ce95e49d16425a8171a7ad1afae5a54 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:47:11 +0200 Subject: [PATCH 1/6] Simplify declaration of werkzeug Map in tests. --- tests/asyncio/test_router.py | 40 ++++++++++++++++-------------------- tests/sync/test_router.py | 36 +++++++++++++++----------------- tests/trio/test_router.py | 38 ++++++++++++++++------------------ 3 files changed, 52 insertions(+), 62 deletions(-) diff --git a/tests/asyncio/test_router.py b/tests/asyncio/test_router.py index ea8e14bf5..c9651e9cc 100644 --- a/tests/asyncio/test_router.py +++ b/tests/asyncio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with route(self.url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: + async with route(url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with route(self.url_map, "localhost", 0, ssl=True) as server: + async with route(url_map, "localhost", 0, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with route(self.url_map, "localhost", 0, server_name="other") as server: + async with route(url_map, "localhost", 0, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -118,7 +116,7 @@ def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -130,7 +128,7 @@ async def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -142,7 +140,7 @@ def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -159,7 +157,7 @@ async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -177,9 +175,7 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with route( - self.url_map, "localhost", 0, create_router=MyRouter - ) as server: + async with route(url_map, "localhost", 0, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index cf04a8689..5c8b4de2d 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) def echo(websocket, count): @@ -48,24 +55,15 @@ def test_router_matches_paths_and_extracts_parameters(self): messages = list(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with connect(get_uri(server) + "/?a=b") as client: self.assertEval(client, "ws.request.path", "/?a=b") def test_redirect(self): """Router redirects connections according to redirect_to.""" - with run_router(self.url_map, server_name="localhost") as server: + with run_router(url_map, server_name="localhost") as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -76,9 +74,7 @@ def test_redirect(self): def test_secure_redirect(self): """Router redirects connections to a wss:// URI when TLS is enabled.""" - with run_router( - self.url_map, server_name="localhost", ssl=SERVER_CONTEXT - ) as server: + with run_router(url_map, server_name="localhost", ssl=SERVER_CONTEXT) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT): self.fail("did not raise") @@ -90,7 +86,7 @@ def test_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - with run_router(self.url_map, ssl=True) as server: + with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): @@ -103,7 +99,7 @@ def test_force_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - with run_router(self.url_map, server_name="other") as server: + with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -114,7 +110,7 @@ def test_force_redirect_server_name(self): def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -129,7 +125,7 @@ def test_process_request_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with connect(get_uri(server) + "/") as client: self.assertEval(client, "ws.process_request_ran", "True") @@ -139,7 +135,7 @@ def test_process_request_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -156,7 +152,7 @@ def handler(self, connection): connection.my_router_ran = True return super().handler(connection) - with run_router(self.url_map, create_router=MyRouter) as server: + with run_router(url_map, create_router=MyRouter) as server: with connect(get_uri(server)) as client: self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/trio/test_router.py b/tests/trio/test_router.py index a85f5bf01..e303a3bd6 100644 --- a/tests/trio/test_router.py +++ b/tests/trio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with run_router(self.url_map, ssl=SERVER_CONTEXT) as server: + async with run_router(url_map, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with run_router(self.url_map, ssl=True) as server: + async with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with run_router(self.url_map, server_name="other") as server: + async with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -117,7 +115,7 @@ async def test_process_request_function_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -127,7 +125,7 @@ async def test_process_request_coroutine_returning_none(self): async def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -137,7 +135,7 @@ async def test_process_request_function_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -152,7 +150,7 @@ async def test_process_request_coroutine_returning_response(self): async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -169,6 +167,6 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with run_router(self.url_map, create_router=MyRouter) as server: + async with run_router(url_map, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") From dbee575ffb7b0bf511b6ef12c16222b57332fa6d Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:09:42 +0200 Subject: [PATCH 2/6] Make implementations of test servers consistent. --- tests/sync/server.py | 12 ++++++++---- tests/sync/test_client.py | 6 +++--- tests/sync/test_server.py | 11 ++++++----- tests/trio/server.py | 30 +++++++++++++++++------------- tests/trio/test_client.py | 20 ++++++++++++-------- tests/trio/test_server.py | 17 +++++++++++------ 6 files changed, 57 insertions(+), 39 deletions(-) diff --git a/tests/sync/server.py b/tests/sync/server.py index 78d1c6745..194829fea 100644 --- a/tests/sync/server.py +++ b/tests/sync/server.py @@ -8,11 +8,15 @@ from websockets.sync.server import serve, unix_serve +def get_host_port(server): + return server.socket.getsockname() + + def get_uri(server, secure=None): if secure is None: secure = isinstance(server.socket, ssl.SSLSocket) # hack protocol = "wss" if secure else "ws" - host, port = server.socket.getsockname() + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -69,9 +73,9 @@ def run_router(url_map, **kwargs): @contextlib.contextmanager def run_unix_server_or_router( - path, unix_serve_or_route, handler_or_url_map, + path, **kwargs, ): with unix_serve_or_route(handler_or_url_map, path, **kwargs) as server: @@ -85,8 +89,8 @@ def run_unix_server_or_router( def run_unix_server(path, handler=handler, **kwargs): - return run_unix_server_or_router(path, unix_serve, handler, **kwargs) + return run_unix_server_or_router(unix_serve, handler, path, **kwargs) def run_unix_router(path, url_map, **kwargs): - return run_unix_server_or_router(path, unix_route, url_map, **kwargs) + return run_unix_server_or_router(unix_route, url_map, path, **kwargs) diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index d4d42c318..1ecb649b9 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -30,7 +30,7 @@ DeprecationTestCase, temp_unix_socket_path, ) -from .server import get_uri, run_server, run_unix_server +from .server import get_host_port, get_uri, run_server, run_unix_server class ClientTests(unittest.TestCase): @@ -43,7 +43,7 @@ def test_connection(self): def test_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -440,7 +440,7 @@ def test_explicit_socks_proxy(self): def test_ignore_proxy_with_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index ee01a5eef..92280ac3f 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -29,6 +29,7 @@ ) from .server import ( EvalShellMixin, + get_host_port, get_uri, handler, run_server, @@ -320,7 +321,7 @@ def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server(open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -334,7 +335,7 @@ def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) @@ -360,7 +361,7 @@ def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets.server", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: sock.send(b"HELO relay.invalid\r\n") # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -502,14 +503,14 @@ def test_connection(self): def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT) as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) diff --git a/tests/trio/server.py b/tests/trio/server.py index 6e9ef417e..7b69e491b 100644 --- a/tests/trio/server.py +++ b/tests/trio/server.py @@ -10,8 +10,8 @@ from websockets.trio.server import serve -def get_host_port(listeners): - for listener in listeners: +def get_host_port(server): + for listener in server.listeners: if listener.socket.family == socket.AF_INET: # pragma: no branch return listener.socket.getsockname() raise AssertionError("expected at least one IPv4 socket") @@ -24,7 +24,7 @@ def get_uri(server, secure=None): for cell in server.handler.__closure__ ) # l33t hack protocol = "wss" if secure else "ws" - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -53,19 +53,23 @@ async def assertEval(self, client, expr, value): self.assertEqual(await client.recv(), value) -kwargs = {"port": 0, "host": "localhost"} - - @contextlib.asynccontextmanager async def run_server_or_route( serve_or_route, handler_or_url_map, - **overrides, + port=0, + host="localhost", + **kwargs, ): - merged_kwargs = {**kwargs, **overrides} async with trio.open_nursery() as nursery: server = await nursery.start( - functools.partial(serve_or_route, handler_or_url_map, **merged_kwargs) + functools.partial( + serve_or_route, + handler_or_url_map, + port, + host=host, + **kwargs, + ) ) try: yield server @@ -76,9 +80,9 @@ async def run_server_or_route( nursery.cancel_scope.cancel() -def run_server(handler=handler, **overrides): - return run_server_or_route(serve, handler, **overrides) +def run_server(handler=handler, **kwargs): + return run_server_or_route(serve, handler, **kwargs) -def run_router(url_map, **overrides): - return run_server_or_route(route, url_map, **overrides) +def run_router(url_map, **kwargs): + return run_server_or_route(route, url_map, **kwargs) diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 6afd35f81..94f925006 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -63,14 +63,14 @@ async def test_connection(self): async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" async with run_server() as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect("ws://overridden/", host=host, port=port) as client: self.assertEqual(client.protocol.state.name, "OPEN") async def test_existing_stream(self): """Client connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -306,7 +306,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "ws://overridden/redirect", host=host, port=port ) as client: @@ -321,7 +321,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) with self.assertRaises(ValueError) as raised: async with connect("ws://overridden/", host=host, port=port): self.fail("did not raise") @@ -341,7 +341,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) with self.assertRaises(ValueError) as raised: # Use a non-existing domain to ensure we connect via sock. async with connect("ws://invalid/redirect", stream=stream): @@ -505,7 +505,11 @@ async def junk(stream): async with trio.open_nursery() as nursery: try: listeners = await nursery.start(trio.serve_tcp, junk, 0) - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() with self.assertRaises(InvalidMessage) as raised: async with connect(f"ws://{host}:{port}"): self.fail("did not raise") @@ -537,7 +541,7 @@ async def test_connection(self): async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" async with run_server(ssl=SERVER_CONTEXT) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "wss://overridden/", host=host, port=port, ssl=CLIENT_CONTEXT ) as client: @@ -721,7 +725,7 @@ async def test_explicit_socks_proxy(self): async def test_ignore_proxy_with_existing_stream(self): """Cli ent connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index 8018175e4..ffd7337c9 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -2,6 +2,7 @@ import hmac import http import logging +import socket import trio @@ -64,7 +65,11 @@ async def test_connection_handler_raises_exception(self): async def test_existing_listeners(self): """Server receives connection using pre-existing listeners.""" listeners = await trio.open_tcp_listeners(0, host="localhost") - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() # Unset the default values of port and host set by run_server. async with run_server(port=None, host=None, listeners=listeners): async with connect(f"ws://{host}:{port}/") as client: # type: ignore @@ -416,7 +421,7 @@ async def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server(open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -433,7 +438,7 @@ async def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose() self.assertExceptionLogged( @@ -459,7 +464,7 @@ async def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.send_all(b"HELO relay.invalid\r\n") try: # Wait for the server to close the connection. @@ -612,7 +617,7 @@ async def test_connection(self): async def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -622,7 +627,7 @@ async def test_timeout_during_tls_handshake(self): async def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose() From b1fe6c14c9710d1d3526e18466d03350168c24da Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 09:19:20 +0200 Subject: [PATCH 3/6] Add tests for connecting without a context manager. --- tests/asyncio/test_client.py | 12 ++++++++++-- tests/sync/test_client.py | 12 ++++++++++-- tests/trio/test_client.py | 5 +++-- 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 6f6070c60..6f0a0c922 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -52,11 +52,19 @@ async def few_redirects(): class ClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to server directly.""" + async with serve(*args) as server: + client = await connect(get_uri(server)) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 1ecb649b9..591accbb0 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -34,11 +34,19 @@ class ClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server and the handshake succeeds.""" + def test_context_manager(self): + """Client connects to server and disconnects automatically.""" with run_server() as server: with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to server directly.""" + with run_server() as server: + client = connect(get_uri(server)) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") def test_existing_socket(self): """Client connects using a pre-existing socket.""" diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 94f925006..dc7c68c97 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -54,11 +54,12 @@ async def few_redirects(): class ClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with run_server() as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" From 620a3fb086e18310d4e7722edf77de59c490c1dc Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 09:20:42 +0200 Subject: [PATCH 4/6] Support overriding host/post in the sync client. --- docs/project/changelog.rst | 6 ++++++ src/websockets/sync/client.py | 10 ++++++---- src/websockets/trio/client.py | 4 ++++ tests/sync/test_client.py | 7 +++++++ 4 files changed, 23 insertions(+), 4 deletions(-) diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index fa2540c11..ab375b481 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -32,6 +32,12 @@ notice. *In development* +New features +............ + +* :func:`~sync.client.connect` can connect to another host and port than those + specified in the URI in the :mod:`threading` implementation. + .. _17.0.1: 17.0.1 diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 509eb9df8..85b9a1a56 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -229,6 +229,10 @@ def connect( connection handling. Any other keyword arguments are passed to :func:`~socket.create_connection`. + For example, you can set ``address`` to a ``(host, port)`` tuple to connect + to a different host and port from those found in ``uri``. This only changes + the destination of the TCP connection. The host name from ``uri`` is still + used in the TLS handshake for secure connections and in the ``Host`` header. Raises: InvalidURI: If ``uri`` isn't a valid WebSocket URI. @@ -331,11 +335,9 @@ def connect( raise AssertionError("parse_proxy returned unsupported proxy") else: # proxy is None + kwargs.setdefault("address", (ws_uri.host, ws_uri.port)) kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection( - (ws_uri.host, ws_uri.port), - **kwargs, - ) + sock = socket.create_connection(**kwargs) sock.settimeout(None) diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 1bdde69c8..9c92d2f67 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -220,6 +220,10 @@ class connect: connection handling. Any other keyword arguments are passed to :func:`~trio.open_tcp_stream`. + For example, you can set ``host`` and ``port`` to connect to a different + host and port from those found in ``uri``. This only changes the destination + of the TCP connection. The host name from ``uri`` is still used in the TLS + handshake for secure connections and in the ``Host`` header. Raises: InvalidURI: If ``uri`` isn't a valid WebSocket URI. diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 591accbb0..70d187f79 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -48,6 +48,13 @@ def test_direct_connection(self): self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") + def test_explicit_host_port(self): + """Client connects using an explicit host / port.""" + with run_server() as server: + address = get_host_port(server) + with connect("ws://overridden/", address=address) as client: + self.assertEqual(client.protocol.state.name, "OPEN") + def test_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: From 020326b63de9ccf3aaa6577c8b828227237bb877 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 18 Aug 2026 15:17:04 +0200 Subject: [PATCH 5/6] Deprecate calling connect() directly. This usage was never shown in examples, but it was clearly described in the documentation and it must exist in the wild. Add a flag for backwards compatibility. While the name implies that it should go away in the long run, it isn't immediately deprecated. Its name provide light deterrence, which is enough for now. --- src/websockets/asyncio/client.py | 11 +++++----- src/websockets/sync/client.py | 35 +++++++++++++++++++++++++++---- src/websockets/sync/connection.py | 9 ++++++++ src/websockets/trio/client.py | 15 ++++++++----- tests/sync/test_client.py | 18 +++++++++++++++- 5 files changed, 72 insertions(+), 16 deletions(-) diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 6c6c11495..f49af72da 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -177,10 +177,9 @@ class connect: """ Connect to the WebSocket server at ``uri``. - This coroutine returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as an asynchronous context manager:: + :func:`connect` should be treated as an asynchronous context manager + yielding a :class:`ClientConnection`, which can then receive and send + messages:: from websockets.asyncio.client import connect @@ -189,8 +188,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 85b9a1a56..6aebf027a 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -65,6 +65,7 @@ def __init__( ) -> None: self.protocol: ClientProtocol self.response_rcvd = threading.Event() + self.pending_legacy_warning = True super().__init__( sock, protocol, @@ -74,6 +75,21 @@ def __init__( max_queue=max_queue, ) + def __enter__(self) -> ClientConnection: + self.pending_legacy_warning = False + return super().__enter__() + + def maybe_raise_legacy_warning(self) -> None: + if self.pending_legacy_warning: + self.pending_legacy_warning = False + warnings.warn( # deprecated in 17.1 + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly", + DeprecationWarning, + stacklevel=3, + ) + def handshake( self, additional_headers: HeadersLike | None = None, @@ -164,10 +180,8 @@ def connect( """ Connect to the WebSocket server at ``uri``. - This function returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as a context manager:: + :func:`connect` should be treated as a context manager yielding a + :class:`ClientConnection`, which can then receive and send messages:: from websockets.sync.client import connect @@ -176,6 +190,13 @@ def connect( The connection is closed automatically when exiting the context. + For backwards compatibility, :func:`connect` may be called directly:: + + websocket = await connect(..., legacy=True) + + In that case, you're responsible for closing the connection with + :meth:`ClientConnection.close` when no longer needed. + Args: uri: URI of the WebSocket server. sock: Preexisting TCP socket. ``sock`` overrides the host and port @@ -253,6 +274,9 @@ def connect( DeprecationWarning, ) + # Backwards compatibility: connect can return a ClientConnection. + legacy = kwargs.pop("legacy", False) + ws_uri = parse_uri(uri) if not ws_uri.secure and ssl is not None: raise ValueError("ssl argument is incompatible with a ws:// URI") @@ -389,6 +413,9 @@ def connect( sock.close() raise + if legacy: + connection.pending_legacy_warning = False + try: connection.handshake( additional_headers, diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997ab..b2b7415cb 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -317,6 +317,7 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: return self.recv_messages.get(timeout, decode) except EOFError: @@ -387,6 +388,7 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: yield from self.recv_messages.get_iter(decode) return @@ -466,6 +468,7 @@ def send( TypeError: If ``message`` doesn't have a supported type. """ + self.maybe_raise_legacy_warning() # Unfragmented message — this case must be handled first because # strings and bytes-like objects are iterable. @@ -591,6 +594,7 @@ def close( reason: WebSocket close reason. """ + self.maybe_raise_legacy_warning() try: # The context manager takes care of waiting for the TCP connection # to terminate after calling a method that sends a close frame. @@ -647,6 +651,7 @@ def ping( the corresponding pong wasn't received yet. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -684,6 +689,7 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -696,6 +702,9 @@ def pong(self, data: DataLike = b"") -> None: # Private methods + def maybe_raise_legacy_warning(self) -> None: + pass # see override in ClientConnection + def process_event(self, event: Event) -> None: """ Process one incoming event. diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 9c92d2f67..cd79f3d69 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -135,8 +135,9 @@ class connect: """ Connect to the WebSocket server at ``uri``. - This coroutine returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. + :func:`connect` is designed to be called as an asynchronous context manager + yielding a :class:`ClientConnection`, which you can then use to receive and + send messages:: :func:`connect` may be used as an asynchronous context manager:: @@ -147,8 +148,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: @@ -162,6 +163,10 @@ class connect: The connection is closed automatically after each iteration of the loop. + :func:`connect` cannot be awaited directly. This is because it runs a task + to manage the connection and Trio doesn't support spawning tasks without a + context that ensures completion. + Args: uri: URI of the WebSocket server. stream: Preexisting TCP stream. ``stream`` overrides the host and port @@ -536,7 +541,7 @@ async def connect(self, nursery: trio.Nursery) -> ClientConnection: # Re-raise exception with an informative error message. raise TimeoutError("timed out during opening handshake") from exc - # Do not define __await__ for... = await nursery.start(connect, ...) + # Do not define __await__ for ... = await nursery.start(connect, ...) # because it doesn't look idiomatic in Trio. # async with connect(...) as ...: ... diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 70d187f79..f65fbd407 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -44,7 +44,7 @@ def test_context_manager(self): def test_direct_connection(self): """Client connects to server directly.""" with run_server() as server: - client = connect(get_uri(server)) + client = connect(get_uri(server), legacy=True) self.addCleanup(client.close) self.assertEqual(client.protocol.state.name, "OPEN") @@ -765,3 +765,19 @@ def test_ssl_context_argument(self): with self.assertDeprecationWarning("ssl_context was renamed to ssl"): with connect(get_uri(server), ssl_context=CLIENT_CONTEXT): pass + + def test_direct_connection_without_legacy_flag(self): + """Client connects to server without legacy=True.""" + with run_server() as server: + client = connect(get_uri(server)) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + # First call of a public API triggers a warning + with self.assertDeprecationWarning( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" + ): + client.ping() + # Later calls don't trigger a warning + client.pong() From 18cfa5dcac39226d0991099ba19e5d0e63cf32d2 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 11 Aug 2026 08:44:45 +0200 Subject: [PATCH 6/6] wip --- docs/project/changelog.rst | 7 + docs/reference/features.rst | 4 +- docs/reference/sync/client.rst | 9 + src/websockets/sync/client.py | 750 +++++++++++++++++++++++------- src/websockets/sync/connection.py | 31 ++ tests/sync/test_client.py | 397 +++++++++++++++- tests/sync/test_router.py | 28 +- 7 files changed, 1026 insertions(+), 200 deletions(-) diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index ab375b481..a7e73ebaf 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -35,6 +35,13 @@ notice. New features ............ +* Added support for reconnecting automatically by using + :func:`~sync.client.reconnect` as an iterator to the :mod:`threading` + implementation. + +* :func:`~sync.client.connect` now follows redirects in the :mod:`threading` + implementation. + * :func:`~sync.client.connect` can connect to another host and port than those specified in the URI in the :mod:`threading` implementation. diff --git a/docs/reference/features.rst b/docs/reference/features.rst index 2bc505b5e..01909bdb1 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -149,7 +149,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Close connection on context exit | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Reconnect automatically | ✅ | ❌ | ✅ | — | ✅ | + | Reconnect automatically | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Configure ``Origin`` header | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ @@ -161,7 +161,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Connect to non-ASCII IRIs | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Follow HTTP redirects | ✅ | ❌ | ✅ | — | ✅ | + | Follow HTTP redirects | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Perform HTTP Basic Authentication | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ diff --git a/docs/reference/sync/client.rst b/docs/reference/sync/client.rst index fdc772ed2..77e95e58e 100644 --- a/docs/reference/sync/client.rst +++ b/docs/reference/sync/client.rst @@ -10,6 +10,15 @@ Opening a connection .. autofunction:: unix_connect +Reconnecting automatically +--------------------------- + +.. autofunction:: reconnect + +.. autofunction:: unix_reconnect + +.. autofunction:: process_exception + Using a connection ------------------ diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 6aebf027a..042129719 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -1,16 +1,28 @@ from __future__ import annotations import logging +import os import socket import ssl as ssl_module import threading +import time +import traceback +import urllib.parse import warnings -from collections.abc import Sequence +from collections.abc import Generator, Iterator, Sequence +from types import TracebackType from typing import Any, Callable, Literal, TypeVar, cast -from ..client import ClientProtocol -from ..datastructures import HeadersLike -from ..exceptions import InvalidProxyMessage, InvalidProxyStatus, ProxyError +from ..asyncio.client import process_exception +from ..client import ClientProtocol, backoff +from ..datastructures import Headers, HeadersLike +from ..exceptions import ( + InvalidProxyMessage, + InvalidProxyStatus, + InvalidStatus, + ProxyError, + SecurityError, +) from ..extensions.base import ClientExtensionFactory from ..extensions.permessage_deflate import enable_client_permessage_deflate from ..headers import validate_subprotocols @@ -24,7 +36,9 @@ from .utils import Deadline -__all__ = ["connect", "unix_connect", "ClientConnection"] +__all__ = ["connect", "unix_connect", "reconnect", "unix_reconnect", "ClientConnection"] + +MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10")) class ClientConnection(Connection): @@ -144,6 +158,500 @@ def recv_events(self) -> None: self.response_rcvd.set() +# This is spelled in lower case because it's exposed as a callable in the API. +class reconnect: + """ + Connect to the WebSocket server at ``uri``, with support for reconnecting. + + This class returns a :class:`ClientConnection` instance, which you can + use to send and receive messages. + + :func:`reconnect` may be used as a context manager for a single + connection attempt:: + + from websockets.sync.client import reconnect + + with reconnect(...) as websocket: + ... + + The connection is closed automatically when exiting the context. + + :func:`reconnect` may be used as an infinite iterator to reconnect + automatically on errors:: + + for websocket in reconnect(...): + try: + ... + except websockets.exceptions.ConnectionClosed: + continue + + If the connection fails with a transient error, it is retried with + exponential backoff. If it fails with a fatal error, the exception is + raised, breaking out of the loop. + + The connection is closed automatically after each iteration of the loop. + + For a single connection attempt without a context manager, see + :func:`connect`. + + Args: + uri: URI of the WebSocket server. + sock: Preexisting TCP socket. ``sock`` overrides the host and port + from ``uri``. You may call :func:`socket.create_connection` to + create a suitable TCP socket. + ssl: Configuration for enabling TLS on the connection. + server_hostname: Host name for the TLS handshake. ``server_hostname`` + overrides the host name from ``uri``. + origin: Value of the ``Origin`` header, for servers that require it. + extensions: List of supported extensions, in order in which they + should be negotiated and run. + subprotocols: List of supported subprotocols, in order of decreasing + preference. + compression: The "permessage-deflate" extension is enabled by default. + Set ``compression`` to :obj:`None` to disable it. See the + :doc:`compression guide <../../topics/compression>` for details. + additional_headers: Arbitrary HTTP headers to add to the handshake + request. + user_agent_header: Value of the ``User-Agent`` request header. + It defaults to ``"Python/x.y.z websockets/X.Y"``. + Setting it to :obj:`None` removes the header. + proxy: If a proxy is configured, it is used by default. Set ``proxy`` + to :obj:`None` to disable the proxy or to the address of a proxy + to override the system configuration. See the :doc:`proxy docs + <../../topics/proxies>` for details. + proxy_ssl: Configuration for enabling TLS on the proxy connection. + proxy_server_hostname: Host name for the TLS handshake with the proxy. + ``proxy_server_hostname`` overrides the host name from ``proxy``. + process_exception: When reconnecting automatically, tell whether an + error is transient or fatal. The default behavior is defined by + :func:`process_exception`. Refer to its documentation for details. + open_timeout: Timeout for opening the connection in seconds. + :obj:`None` disables the timeout. + ping_interval: Interval between keepalive pings in seconds. + :obj:`None` disables keepalive. + ping_timeout: Timeout for keepalive pings in seconds. + :obj:`None` disables timeouts. + close_timeout: Timeout for closing the connection in seconds. + :obj:`None` disables the timeout. + reconnect_delays: Delays in seconds between reconnection attempts. + Default is exponential backoff with 5s jitter, capped at 60s. + max_size: Maximum size of incoming messages in bytes. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. + max_queue: High-water mark of the buffer where frames are received. + It defaults to 16 frames. The low-water mark defaults to ``max_queue + // 4``. You may pass a ``(high, low)`` tuple to set the high-water + and low-water marks. If you want to disable flow control entirely, + you may set it to ``None``, although that's a bad idea. + logger: Logger for this client. + It defaults to ``logging.getLogger("websockets.client")``. + See the :doc:`logging guide <../../topics/logging>` for details. + create_connection: Factory for the :class:`ClientConnection` managing + the connection. Set it to a wrapper or a subclass to customize + connection handling. + + Any other keyword arguments are passed to :func:`~socket.create_connection`. + For example, you can set ``address`` to a ``(host, port)`` tuple to connect + to a different host and port from those found in ``uri``. This only changes + the destination of the TCP connection. The host name from ``uri`` is still + used in the TLS handshake for secure connections and in the ``Host`` header. + + Raises: + InvalidURI: If ``uri`` isn't a valid WebSocket URI. + InvalidProxy: If ``proxy`` isn't a valid proxy. + OSError: If the TCP connection fails. + InvalidHandshake: If the opening handshake fails. + TimeoutError: If the opening handshake times out. + + """ + + def __init__( + self, + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = None, + ssl: ssl_module.SSLContext | None = None, + server_hostname: str | None = None, + # WebSocket + origin: Origin | None = None, + extensions: Sequence[ClientExtensionFactory] | None = None, + subprotocols: Sequence[Subprotocol] | None = None, + compression: str | None = "deflate", + # HTTP + additional_headers: HeadersLike | None = None, + user_agent_header: str | None = USER_AGENT, + proxy: str | Literal[True] | None = True, + proxy_ssl: ssl_module.SSLContext | None = None, + proxy_server_hostname: str | None = None, + process_exception: Callable[[Exception], Exception | None] = process_exception, + # Timeouts + open_timeout: float | None = 10, + ping_interval: float | None = 20, + ping_timeout: float | None = 20, + close_timeout: float | None = 10, + reconnect_delays: Callable[[], Generator[float]] = backoff, + # Limits + max_size: int | None | tuple[int | None, int | None] = 2**20, + max_queue: int | None | tuple[int | None, int | None] = 16, + # Logging + logger: LoggerLike | None = None, + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = None, + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, + ) -> None: + # Backwards compatibility: ssl used to be called ssl_context. + if ssl is None and "ssl_context" in kwargs: + ssl = kwargs.pop("ssl_context") + warnings.warn( # deprecated in 13.0 - 2024-08-20 + "ssl_context was renamed to ssl", + DeprecationWarning, + ) + + self.uri = uri + self.ws_uri = parse_uri(uri) + if not self.ws_uri.secure and ssl is not None: + raise ValueError("ssl argument is incompatible with a ws:// URI") + + if subprotocols is not None: + validate_subprotocols(subprotocols) + + if compression == "deflate": + extensions = enable_client_permessage_deflate(extensions) + elif compression is not None: + raise ValueError(f"unsupported compression: {compression}") + + if logger is None: + logger = logging.getLogger("websockets.client") + + if create_connection is None: + create_connection = ClientConnection + + # Private APIs for unix_connect() + unix: bool = kwargs.pop("unix", False) + path: str | None = kwargs.pop("path", None) + + if unix: + if path is None and sock is None: + raise ValueError("missing path argument") + elif path is not None and sock is not None: + raise ValueError("path is incompatible with sock") + kwargs["unix"] = True + kwargs["path"] = path + + self.sock = sock + self.ssl = ssl + self.server_hostname = server_hostname + self.additional_headers = additional_headers + self.user_agent_header = user_agent_header + self.proxy = proxy + self.proxy_ssl = proxy_ssl + self.proxy_server_hostname = proxy_server_hostname + self.process_exception = process_exception + self.open_timeout = open_timeout + self.reconnect_delays = reconnect_delays + self.logger = logger + self.create_connection = create_connection + self.open_socket_kwargs = kwargs + self.protocol_kwargs = dict( + origin=origin, + extensions=extensions, + subprotocols=subprotocols, + max_size=max_size, + logger=logger, + ) + self.connection_kwargs = dict( + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + max_queue=max_queue, + ) + + def open_socket(self, deadline: Deadline) -> socket.socket: + """Open a TCP or Unix connection to the server, possibly through a proxy.""" + kwargs = self.open_socket_kwargs.copy() + + proxy = self.proxy + if kwargs.get("unix", False): + proxy = None + if proxy is True: + proxy = get_proxy(self.ws_uri) + + if kwargs.pop("unix", False): + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.settimeout(deadline.timeout()) + sock.connect(kwargs.pop("path")) + except Exception: + sock.close() + raise + + elif proxy is not None: + proxy_parsed = parse_proxy(proxy) + + if proxy_parsed.scheme[:5] == "socks": + sock = connect_socks_proxy( + proxy_parsed, + self.ws_uri, + deadline, + # websockets is consistent with the socket module while + # python_socks is consistent across implementations. + local_addr=kwargs.pop("source_address", None), + ) + + elif proxy_parsed.scheme[:4] == "http": + if proxy_parsed.scheme != "https" and self.proxy_ssl is not None: + raise ValueError( + "proxy_ssl argument is incompatible with an http:// proxy" + ) + sock = connect_http_proxy( + proxy_parsed, + self.ws_uri, + deadline, + user_agent_header=self.user_agent_header, + ssl=self.proxy_ssl, + server_hostname=self.proxy_server_hostname, + **kwargs, + ) + + else: + raise AssertionError("parse_proxy returned unsupported proxy") + + else: # proxy is None + kwargs.setdefault("address", (self.ws_uri.host, self.ws_uri.port)) + kwargs.setdefault("timeout", deadline.timeout()) + sock = socket.create_connection(**kwargs) + + sock.settimeout(None) + return sock + + def enable_tls(self, sock: socket.socket, deadline: Deadline) -> socket.socket: + """Enable TLS on the connection.""" + if self.ssl is None: + ssl = ssl_module.create_default_context() + else: + ssl = self.ssl + if self.server_hostname is None: + server_hostname = self.ws_uri.host + else: + server_hostname = self.server_hostname + sock.settimeout(deadline.timeout()) + if self.proxy_ssl is None: + sock = ssl.wrap_socket(sock, server_hostname=server_hostname) + else: + sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) + # Let's pretend that sock is a socket, even though it isn't. + sock = cast(socket.socket, sock_2) + sock.settimeout(None) + return sock + + def open_connection(self, deadline: Deadline) -> ClientConnection: + """Create a WebSocket connection.""" + # TCP connection is already established. + if self.sock is not None: + sock = self.sock + else: + sock = self.open_socket(deadline) + + try: + # Disable Nagle algorithm + + if not self.open_socket_kwargs.get("unix", False): + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) + + # Initialize TLS wrapper and perform TLS handshake + + if self.ws_uri.secure: + sock = self.enable_tls(sock, deadline) + + # Initialize WebSocket protocol + + protocol = ClientProtocol( + self.ws_uri, + **self.protocol_kwargs, # type: ignore + ) + + # Initialize WebSocket connection + + # self.create_connection defaults to ClientConnection. + connection = self.create_connection( + sock, + protocol, + **self.connection_kwargs, # type: ignore + ) + except Exception: + sock.close() + raise + + try: + connection.handshake( + self.additional_headers, + self.user_agent_header, + deadline.timeout(), + ) + except Exception: + connection.close_socket() + connection.recv_events_thread.join() + raise + + return connection + + def process_redirect(self, exc: Exception) -> Exception | str: + """ + Determine whether a connection error is a redirect that can be followed. + + Return the new URI if it's a valid redirect. Else, return an exception. + + """ + if not ( + isinstance(exc, InvalidStatus) + and exc.response.status_code + in [ + 300, # Multiple Choices + 301, # Moved Permanently + 302, # Found + 303, # See Other + 307, # Temporary Redirect + 308, # Permanent Redirect + ] + and "Location" in exc.response.headers + ): + return exc + + old_ws_uri = self.ws_uri + new_uri = urllib.parse.urljoin(self.uri, exc.response.headers["Location"]) + new_ws_uri = parse_uri(new_uri) + + # If connect() received a socket, it is closed and cannot be reused. + if self.sock is not None: + return ValueError( + f"cannot follow redirect to {new_uri} with a preexisting socket" + ) + + # TLS downgrade is forbidden. + if old_ws_uri.secure and not new_ws_uri.secure: + return SecurityError(f"cannot follow redirect to non-secure URI {new_uri}") + + # Apply restrictions to cross-origin redirects. + if ( + old_ws_uri.secure != new_ws_uri.secure + or old_ws_uri.host != new_ws_uri.host + or old_ws_uri.port != new_ws_uri.port + ): + # Cross-origin redirects on Unix sockets don't quite make sense. + if self.open_socket_kwargs.get("unix", False): + return ValueError( + f"cannot follow cross-origin redirect to {new_uri} " + f"with a Unix socket" + ) + # Cross-origin redirects when host and port are overridden are ill-defined. + if self.open_socket_kwargs.get("address") is not None: + return ValueError( + f"cannot follow cross-origin redirect to {new_uri} " + f"with an explicit host or port" + ) + + # Strip credentials to avoid leaking them to a different origin. + if self.additional_headers is not None: + self.additional_headers = Headers( + ( + (key, value) + for key, value in Headers(self.additional_headers).raw_items() + if key.lower() + not in ["authorization", "cookie", "proxy-authorization"] + ) + ) + + return new_uri + + def connect(self) -> ClientConnection: + """Connect to a WebSocket server, following redirects.""" + deadline = Deadline(self.open_timeout) + for _ in range(MAX_REDIRECTS): + try: + connection = self.open_connection(deadline) + except Exception as exc: + exc_or_uri = self.process_redirect(exc) + if isinstance(exc_or_uri, Exception): + # Response isn't a valid redirect; raise the exception. + if exc_or_uri is exc: + raise + else: + raise exc_or_uri from exc + else: + # Response is a valid redirect; follow it. + self.uri = exc_or_uri + self.ws_uri = parse_uri(exc_or_uri) + continue + else: + connection.start_keepalive() + return connection + else: + raise SecurityError(f"more than {MAX_REDIRECTS} redirects") + + # with reconnect(...) as ...: ... + + def __enter__(self) -> ClientConnection: + if hasattr(self, "connection"): + raise RuntimeError("reconnect() isn't reentrant") + self.connection = self.connect() + return self.connection + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_traceback: TracebackType | None, + ) -> None: + try: + self.connection.close() + finally: + del self.connection + + # for ... in reconnect(...): ... + + def __iter__(self) -> Iterator[ClientConnection]: + delays: Generator[float] | None = None + while True: + try: + with self as connection: + yield connection + except Exception as exc: + # Determine whether the exception is retryable or fatal. + # The API of process_exception is "return an exception or None"; + # "raise an exception" is also supported because it's a frequent + # mistake. It isn't documented in order to keep the API simple. + try: + new_exc = self.process_exception(exc) + except Exception as raised_exc: + new_exc = raised_exc + + # The connection failed with a fatal error. + # Raise the exception and exit the loop. + if new_exc is exc: + raise + if new_exc is not None: + raise new_exc from exc + + # The connection failed with a retryable error. + # Start or continue backoff and reconnect. + if delays is None: + delays = self.reconnect_delays() + delay = next(delays) + self.logger.info( + "connect failed; reconnecting in %.1f seconds: %s", + delay, + traceback.format_exception_only(exc)[0].strip(), + ) + time.sleep(delay) + + else: + # The connection succeeded. Reset backoff. + delays = None + + def connect( uri: str, *, @@ -174,6 +682,8 @@ def connect( logger: LoggerLike | None = None, # Escape hatch for advanced customization create_connection: type[ClientConnection] | None = None, + # Deprecation + legacy: bool = False, # Other keyword arguments are passed to socket.create_connection **kwargs: Any, ) -> ClientConnection: @@ -190,12 +700,19 @@ def connect( The connection is closed automatically when exiting the context. + Use :func:`reconnect` to reconnect automatically on errors. + For backwards compatibility, :func:`connect` may be called directly:: websocket = await connect(..., legacy=True) In that case, you're responsible for closing the connection with - :meth:`ClientConnection.close` when no longer needed. + :meth:`ClientConnection.close` wh`en no longer needed. + + .. When the ``legacy`` flag is enabled, :func:`connect` returns directly a + .. :class:`ClientConnection` and iterating that connection yields messages. + .. This is different from the default behavior of returning an object that + .. can be iterated to reconnect automatically. Args: uri: URI of the WebSocket server. @@ -248,6 +765,8 @@ def connect( create_connection: Factory for the :class:`ClientConnection` managing the connection. Set it to a wrapper or a subclass to customize connection handling. + legacy: Set this to :obj:`True` to use :func:`connect` outside a + context manager without triggering a deprecation warning. Any other keyword arguments are passed to :func:`~socket.create_connection`. For example, you can set ``address`` to a ``(host, port)`` tuple to connect @@ -263,177 +782,43 @@ def connect( TimeoutError: If the opening handshake times out. """ - - # Process parameters - - # Backwards compatibility: ssl used to be called ssl_context. - if ssl is None and "ssl_context" in kwargs: - ssl = kwargs.pop("ssl_context") - warnings.warn( # deprecated in 13.0 - 2024-08-20 - "ssl_context was renamed to ssl", - DeprecationWarning, - ) - # Backwards compatibility: connect can return a ClientConnection. legacy = kwargs.pop("legacy", False) - ws_uri = parse_uri(uri) - if not ws_uri.secure and ssl is not None: - raise ValueError("ssl argument is incompatible with a ws:// URI") - - if subprotocols is not None: - validate_subprotocols(subprotocols) - - if compression == "deflate": - extensions = enable_client_permessage_deflate(extensions) - elif compression is not None: - raise ValueError(f"unsupported compression: {compression}") - - if logger is None: - logger = logging.getLogger("websockets.client") - - if create_connection is None: - create_connection = ClientConnection - - # Private APIs for unix_connect() - unix: bool = kwargs.pop("unix", False) - path: str | None = kwargs.pop("path", None) - - if unix: - if path is None and sock is None: - raise ValueError("missing path argument") - elif path is not None and sock is not None: - raise ValueError("path is incompatible with sock") - - if unix: - proxy = None - if sock is not None: - proxy = None - if proxy is True: - proxy = get_proxy(ws_uri) - - # Calculate timeouts on the TCP, TLS, and WebSocket handshakes. - # The TCP and TLS timeouts must be set on the socket, then removed - # to avoid conflicting with the WebSocket timeout in handshake(). - deadline = Deadline(open_timeout) - - try: - # Connect socket - - if sock is None: - if unix: - sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - sock.settimeout(deadline.timeout()) - assert path is not None # mypy cannot figure this out - sock.connect(path) - - elif proxy is not None: - proxy_parsed = parse_proxy(proxy) - - if proxy_parsed.scheme[:5] == "socks": - sock = connect_socks_proxy( - proxy_parsed, - ws_uri, - deadline, - # websockets is consistent with the socket module while - # python_socks is consistent across implementations. - local_addr=kwargs.pop("source_address", None), - ) - - elif proxy_parsed.scheme[:4] == "http": - if proxy_parsed.scheme != "https" and proxy_ssl is not None: - raise ValueError( - "proxy_ssl argument is incompatible with an http:// proxy" - ) - sock = connect_http_proxy( - proxy_parsed, - ws_uri, - deadline, - user_agent_header=user_agent_header, - ssl=proxy_ssl, - server_hostname=proxy_server_hostname, - **kwargs, - ) - - else: - raise AssertionError("parse_proxy returned unsupported proxy") - - else: # proxy is None - kwargs.setdefault("address", (ws_uri.host, ws_uri.port)) - kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection(**kwargs) - - sock.settimeout(None) - - # Disable Nagle algorithm - - if not unix: - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) - - # Initialize TLS wrapper and perform TLS handshake - - if ws_uri.secure: - if ssl is None: - ssl = ssl_module.create_default_context() - if server_hostname is None: - server_hostname = ws_uri.host - sock.settimeout(deadline.timeout()) - if proxy_ssl is None: - sock = ssl.wrap_socket(sock, server_hostname=server_hostname) - else: - sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) - # Let's pretend that sock is a socket, even though it isn't. - sock = cast(socket.socket, sock_2) - sock.settimeout(None) - - # Initialize WebSocket protocol - - protocol = ClientProtocol( - ws_uri, - origin=origin, - extensions=extensions, - subprotocols=subprotocols, - max_size=max_size, - logger=logger, - ) - - # Initialize WebSocket connection - - # create_connection defaults to ClientConnection. - connection = create_connection( - sock, - protocol, - ping_interval=ping_interval, - ping_timeout=ping_timeout, - close_timeout=close_timeout, - max_queue=max_queue, - ) - except Exception: - if sock is not None: - sock.close() - raise - + connection = reconnect( + uri, + sock=sock, + ssl=ssl, + server_hostname=server_hostname, + origin=origin, + extensions=extensions, + subprotocols=subprotocols, + compression=compression, + additional_headers=additional_headers, + user_agent_header=user_agent_header, + proxy=proxy, + proxy_ssl=proxy_ssl, + proxy_server_hostname=proxy_server_hostname, + open_timeout=open_timeout, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + max_size=max_size, + max_queue=max_queue, + logger=logger, + create_connection=create_connection, + **kwargs, + ).connect() if legacy: connection.pending_legacy_warning = False - - try: - connection.handshake( - additional_headers, - user_agent_header, - deadline.timeout(), - ) - except Exception: - connection.close_socket() - connection.recv_events_thread.join() - raise - - connection.start_keepalive() return connection def unix_connect( path: str | None = None, uri: str | None = None, + *, + legacy: bool = False, **kwargs: Any, ) -> ClientConnection: """ @@ -458,7 +843,38 @@ def unix_connect( uri = "ws://localhost/" else: uri = "wss://localhost/" - return connect(uri=uri, unix=True, path=path, **kwargs) + return connect(uri=uri, unix=True, path=path, legacy=legacy, **kwargs) + + +def unix_reconnect( + path: str | None = None, + uri: str | None = None, + **kwargs: Any, +) -> reconnect: + """ + Connect to a WebSocket server listening on a Unix socket, with support + for reconnecting. + + This function accepts the same keyword arguments as :func:`reconnect`. + + It's only available on Unix. + + It's mainly useful for debugging servers listening on Unix sockets. + + Args: + path: File system path to the Unix socket. + uri: URI of the WebSocket server. ``uri`` defaults to + ``ws://localhost/`` or, when a ``ssl`` is provided, to + ``wss://localhost/``. + + """ + if uri is None: + # Backwards compatibility: ssl used to be called ssl_context. + if kwargs.get("ssl") is None and kwargs.get("ssl_context") is None: + uri = "ws://localhost/" + else: + uri = "wss://localhost/" + return reconnect(uri=uri, unix=True, path=path, **kwargs) try: diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index b2b7415cb..1cefc9c46 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -317,7 +317,11 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data :meth:`recv_streaming` concurrently. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) try: return self.recv_messages.get(timeout, decode) except EOFError: @@ -388,7 +392,11 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: :meth:`recv_streaming` concurrently. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) try: yield from self.recv_messages.get_iter(decode) return @@ -468,7 +476,11 @@ def send( TypeError: If ``message`` doesn't have a supported type. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) # Unfragmented message — this case must be handled first because # strings and bytes-like objects are iterable. @@ -594,7 +606,11 @@ def close( reason: WebSocket close reason. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) try: # The context manager takes care of waiting for the TCP connection # to terminate after calling a method that sends a close frame. @@ -651,7 +667,11 @@ def ping( the corresponding pong wasn't received yet. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -689,7 +709,11 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ +<<<<<<< HEAD self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -702,8 +726,15 @@ def pong(self, data: DataLike = b"") -> None: # Private methods +<<<<<<< HEAD def maybe_raise_legacy_warning(self) -> None: pass # see override in ClientConnection +======= + def raise_legacy_warning(self) -> None: + # No-op here. websockets.sync.client.ClientConnection overrides this + # method to warn when connect() is used outside a context manager. + pass +>>>>>>> e97cd2f (wip) def process_event(self, event: Event) -> None: """ diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index f65fbd407..a79e9bcc0 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -1,3 +1,4 @@ +import contextlib import http import logging import os @@ -8,8 +9,10 @@ import threading import time import unittest +import warnings from unittest.mock import patch +from websockets.client import backoff from websockets.exceptions import ( InvalidHandshake, InvalidMessage, @@ -18,6 +21,7 @@ InvalidStatus, InvalidURI, ProxyError, + SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate from websockets.sync.client import * @@ -33,6 +37,28 @@ from .server import get_host_port, get_uri, run_server, run_unix_server +def short_backoff(): + defaults = backoff.__defaults__ + yield from backoff( + defaults[0] * MS, + defaults[1] * MS, + defaults[2] * MS, + defaults[3], + ) + + +@contextlib.contextmanager +def few_redirects(): + from websockets.sync import client + + max_redirects = client.MAX_REDIRECTS + client.MAX_REDIRECTS = 2 + try: + yield + finally: + client.MAX_REDIRECTS = max_redirects + + class ClientTests(unittest.TestCase): def test_context_manager(self): """Client connects to server and disconnects automatically.""" @@ -51,7 +77,7 @@ def test_direct_connection(self): def test_explicit_host_port(self): """Client connects using an explicit host / port.""" with run_server() as server: - address = get_host_port(server) + address = get_host_port(server.socket) with connect("ws://overridden/", address=address) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -142,6 +168,248 @@ def create_connection(*args, **kwargs): ) as client: self.assertTrue(client.create_connection_ran) + def test_reconnect(self): + """Client reconnects to server.""" + iterations = 0 + successful = 0 + + def process_request(connection, request): + nonlocal iterations + iterations += 1 + # Retriable errors + if iterations == 1: + time.sleep(3 * MS) + elif iterations == 2: + connection.socket.close() + elif iterations == 3: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + # Fatal error + elif iterations == 6: + return connection.respond(http.HTTPStatus.PAYMENT_REQUIRED, "💸") + + with run_server(process_request=process_request) as server: + with self.assertRaises(InvalidStatus) as raised: + for client in reconnect( + get_uri(server), + open_timeout=3 * MS, + reconnect_delays=short_backoff, + ): + self.assertEqual(client.protocol.state.name, "OPEN") + successful += 1 + + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 402", + ) + self.assertEqual(iterations, 6) + self.assertEqual(successful, 2) + + def test_reconnect_with_custom_process_exception(self): + """Client runs process_exception to tell if errors are retryable or fatal.""" + iteration = 0 + + def process_request(connection, request): + nonlocal iteration + iteration += 1 + if iteration == 1: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus): + if 500 <= exc.response.status_code < 600: + return None + if exc.response.status_code == 418: + return Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual(iteration, 2) + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + + def test_reconnect_with_custom_process_exception_raising_exception(self): + """Client supports raising an exception in process_exception.""" + + def process_request(connection, request): + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus) and exc.response.status_code == 418: + raise Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + + def test_redirect(self): + """Client follows redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect(get_uri(server) + "/redirect") as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect(get_uri(server)): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + @few_redirects() + def test_redirect_limit(self): + """Client stops following redirects after limit is reached.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = request.path + return response + + with run_server(process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + with connect(get_uri(server)): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "more than 2 redirects", + ) + + def test_redirect_with_explicit_host_port(self): + """Client follows redirect with an explicit host / port.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server.socket) + with connect("ws://overridden/redirect", address=address) as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect_with_explicit_host_port(self): + """Client doesn't follow cross-origin redirect with an explicit host / port.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "ws://other/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server.socket) + with self.assertRaises(ValueError) as raised: + with connect("ws://overridden/", address=address): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow cross-origin redirect to ws://other/ " + "with an explicit host or port", + ) + + def test_redirect_with_existing_socket(self): + """Client doesn't follow redirect when using a pre-existing socket.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with socket.create_connection(get_host_port(server.socket)) as sock: + with self.assertRaises(ValueError) as raised: + # Use a non-existing domain to ensure we connect via sock. + with connect("ws://invalid/redirect", sock=sock): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow redirect to ws://invalid/ with a preexisting socket", + ) + + def test_cross_origin_redirect_strips_credentials(self): + """Client strips credentials when following a cross-origin redirect.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect( + get_uri(server), + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertNotIn("Authorization", client.request.headers) + self.assertNotIn("Cookie", client.request.headers) + self.assertNotIn("Proxy-Authorization", client.request.headers) + self.assertIn("X-Custom", client.request.headers) + + def test_same_origin_redirect_preserves_credentials(self): + """Client preserves credentials when following a same-origin redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect( + get_uri(server) + "/redirect", + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertIn("Authorization", client.request.headers) + self.assertIn("Cookie", client.request.headers) + self.assertIn("Proxy-Authorization", client.request.headers) + def test_invalid_uri(self): """Client receives an invalid URI.""" with self.assertRaises(InvalidURI): @@ -361,6 +629,40 @@ def test_reject_invalid_server_hostname(self): str(raised.exception), ) + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with run_server(ssl=SERVER_CONTEXT) as other_server: + with connect(get_uri(server), ssl=CLIENT_CONTEXT): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + def test_redirect_to_insecure_uri(self): + """Client doesn't follow redirect from secure URI to non-secure URI.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = insecure_uri + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + secure_uri = get_uri(server) + insecure_uri = secure_uri.replace("wss://", "ws://") + with connect(secure_uri, ssl=CLIENT_CONTEXT): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + f"cannot follow redirect to non-secure URI {insecure_uri}", + ) + @unittest.skipUnless("mitmproxy" in sys.modules, "mitmproxy not installed") class SocksProxyClientTests(ProxyMixin, unittest.TestCase): @@ -699,25 +1001,17 @@ def test_ssl_without_secure_uri(self): def test_proxy_ssl_without_https_proxy(self): """Client rejects proxy_ssl when proxy isn't HTTPS.""" with self.assertRaises(ValueError) as raised: - connect( + with connect( "ws://localhost/", proxy="http://localhost:8080", proxy_ssl=CLIENT_CONTEXT, - ) + ): + self.fail("did not raise") self.assertEqual( str(raised.exception), "proxy_ssl argument is incompatible with an http:// proxy", ) - def test_unix_without_path_or_sock(self): - """Unix client requires path when sock isn't provided.""" - with self.assertRaises(ValueError) as raised: - unix_connect() - self.assertEqual( - str(raised.exception), - "missing path argument", - ) - def test_unsupported_proxy(self): """Client rejects unsupported proxy.""" with self.assertRaises(InvalidProxy) as raised: @@ -728,6 +1022,15 @@ def test_unsupported_proxy(self): "other://localhost:58080 isn't a valid proxy: scheme other isn't supported", ) + def test_unix_without_path_or_sock(self): + """Unix client requires path when sock isn't provided.""" + with self.assertRaises(ValueError) as raised: + unix_connect() + self.assertEqual( + str(raised.exception), + "missing path argument", + ) + def test_unix_with_path_and_sock(self): """Unix client rejects path when sock is provided.""" sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -757,6 +1060,19 @@ def test_unsupported_compression(self): "unsupported compression: False", ) + def test_reentrancy(self): + """reconnect() isn't reentrant.""" + with run_server() as server: + connecter = reconnect(get_uri(server)) + with connecter: + with self.assertRaises(RuntimeError) as raised: + with connecter: + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "reconnect() isn't reentrant", + ) + class BackwardsCompatibilityTests(DeprecationTestCase): def test_ssl_context_argument(self): @@ -781,3 +1097,60 @@ def test_direct_connection_without_legacy_flag(self): client.ping() # Later calls don't trigger a warning client.pong() + +LEGACY_WARNING = ( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" +) + + +class LegacyConnectTests(DeprecationTestCase): + def test_legacy_warns_on_first_use(self): + """connect() without a context manager warns on first use.""" + with run_server() as server: + client = connect(get_uri(server)) + with self.assertDeprecationWarning(LEGACY_WARNING): + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + + def test_legacy_warns_on_message_iteration(self): + """connect() without a context manager warns when iterating messages.""" + with run_server() as server: + client = connect(get_uri(server) + "/no-op") + with self.assertDeprecationWarning(LEGACY_WARNING): + messages = list(client) + self.assertEqual(messages, []) + + def test_legacy_warns_once(self): + """connect() without a context manager warns only once.""" + with run_server() as server: + client = connect(get_uri(server)) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + self.assertEqual(len(recorded_warnings), 1) + + def test_legacy_true_doesnt_warn(self): + """connect(legacy=True) doesn't warn.""" + with run_server() as server: + client = connect(get_uri(server), legacy=True) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + self.assertEqual(recorded_warnings, []) + + def test_context_manager_doesnt_warn(self): + """connect() doesn't warn when used as a context manager.""" + with run_server() as server: + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + with connect(get_uri(server)) as client: + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + self.assertEqual(recorded_warnings, []) diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index 5c8b4de2d..4d2b2ad99 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -63,27 +63,17 @@ def test_route_with_query_string(self): def test_redirect(self): """Router redirects connections according to redirect_to.""" - with run_router(url_map, server_name="localhost") as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r"): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "ws://localhost/", - ) + with run_router(url_map) as server: + with connect(get_uri(server) + "/r") as client: + self.assertEval(client, "ws.request.path", "/") def test_secure_redirect(self): - """Router redirects connections to a wss:// URI when TLS is enabled.""" - with run_router(url_map, server_name="localhost", ssl=SERVER_CONTEXT) as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "wss://localhost/", - ) + """Router redirects connections according to redirect_to when TLS is enabled.""" + with run_router(url_map, ssl=SERVER_CONTEXT) as server: + with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: + self.assertEval(client, "ws.request.path", "/") - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" with run_router(url_map, ssl=True) as server: @@ -96,7 +86,7 @@ def test_force_secure_redirect(self): redirect_uri + "/", ) - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" with run_router(url_map, server_name="other") as server: