diff --git a/tests/resources/jupyter_server_config.py b/tests/resources/jupyter_server_config.py index 97136d2f..d8dee097 100644 --- a/tests/resources/jupyter_server_config.py +++ b/tests/resources/jupyter_server_config.py @@ -127,6 +127,15 @@ def my_env(): "rewrite_response": [cats_only, dog_to_cat], }, "python-proxyto54321-no-command": {"port": 54321}, + "python-rawsocket-tcp": { + "command": [sys.executable, "./tests/resources/rawsocket.py", "{port}"], + "raw_socket_proxy": True + }, + "python-rawsocket-unix": { + "command": [sys.executable, "./tests/resources/rawsocket.py", "{unix_socket}"], + "unix_socket": True, + "raw_socket_proxy": True + }, } c.ServerProxy.non_service_rewrite_response = hello_to_foo diff --git a/tests/resources/rawsocket.py b/tests/resources/rawsocket.py new file mode 100644 index 00000000..d95c5cd8 --- /dev/null +++ b/tests/resources/rawsocket.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python + +import os +import socket +import sys + +if len(sys.argv) != 2: + print(f"Usage: {sys.argv[0]} TCPPORT|SOCKPATH") + sys.exit(1) +where = sys.argv[1] +try: + port = int(where) + family = socket.AF_INET + addr = ('localhost', port) +except ValueError: + family = socket.AF_UNIX + addr = where + +with socket.create_server(addr, family=family, reuse_port=True) as serv: + while True: + # only handle a single connection at a time + sock, caddr = serv.accept() + while True: + s = sock.recv(1024) + if not s: + break + sock.send(s.swapcase()) + sock.close() diff --git a/tests/test_proxies.py b/tests/test_proxies.py index 8c2de5a0..4c0bb80f 100644 --- a/tests/test_proxies.py +++ b/tests/test_proxies.py @@ -469,3 +469,25 @@ def test_callable_environment_formatting( PORT, TOKEN = a_server_port_and_token r = request_get(PORT, "/python-http-callable-env/test", TOKEN) assert r.code == 200 + + +@pytest.mark.parametrize("rawsocket_type", [ + "tcp", + pytest.param( + "unix", + marks=pytest.mark.skipif( + sys.platform == "win32", reason="Unix socket not supported on Windows" + ), + ), +]) +async def test_server_proxy_rawsocket( + rawsocket_type: str, + a_server_port_and_token: Tuple[int, str] +) -> None: + PORT, TOKEN = a_server_port_and_token + url = f"ws://{LOCALHOST}:{PORT}/python-rawsocket-{rawsocket_type}/?token={TOKEN}" + conn = await websocket_connect(url) + for msg in [b"Hello,", b"world!"]: + await conn.write_message(msg) + res = await conn.read_message() + assert res == msg.swapcase()