chore: import upstream snapshot with attribution

python/ray/serve/tests/test_websockets.py

@@ -0,0 +1,158 @@
+import asyncio
+import sys
+
+import httpx
+import pytest
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from starlette.responses import StreamingResponse
+from websockets.exceptions import ConnectionClosed
+from websockets.sync.client import connect
+
+import ray
+from ray import serve
+from ray._common.test_utils import SignalActor
+from ray.serve._private.test_utils import get_application_url
+
+
+@pytest.mark.parametrize("route_prefix", [None, "/prefix"])
+def test_send_recv_text_and_binary(serve_instance, route_prefix: str):
+    app = FastAPI()
+
+    @serve.deployment
+    @serve.ingress(app)
+    class WebSocketServer:
+        @app.websocket("/")
+        async def ws_handler(self, ws: WebSocket):
+            await ws.accept()
+
+            text = await ws.receive_text()
+            await ws.send_text(text)
+
+            bytes = await ws.receive_bytes()
+            await ws.send_bytes(bytes)
+
+    serve.run(WebSocketServer.bind(), route_prefix=route_prefix or "/")
+
+    msg = "Hello world!"
+    url = f"{get_application_url(is_websocket=True, use_localhost=True)}/"
+
+    with connect(url) as websocket:
+        websocket.send(msg)
+        assert websocket.recv() == msg
+
+        websocket.send(msg.encode("utf-8"))
+        assert websocket.recv().decode("utf-8") == msg
+
+
+def test_client_disconnect(serve_instance):
+    app = FastAPI()
+
+    @serve.deployment
+    @serve.ingress(app)
+    class WebSocketServer:
+        def __init__(self):
+            self._disconnected = asyncio.Event()
+
+        async def wait_for_disconnect(self):
+            await self._disconnected.wait()
+
+        @app.websocket("/")
+        async def ws_handler(self, ws: WebSocket):
+            await ws.accept()
+
+            try:
+                await ws.receive_text()
+            except WebSocketDisconnect:
+                print("Client disconnected.")
+                self._disconnected.set()
+
+    h = serve.run(WebSocketServer.bind())
+    wait_response = h.wait_for_disconnect.remote()
+    url = f"{get_application_url(is_websocket=True)}/"
+
+    with connect(url):
+        print("Client connected.")
+
+    wait_response.result()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Hanging on Windows.")
+def test_server_disconnect(serve_instance):
+    """Test that server can properly close WebSocket connections."""
+    app = FastAPI()
+
+    @serve.deployment
+    @serve.ingress(app)
+    class WebSocketServer:
+        @app.websocket("/")
+        async def ws_handler(self, ws: WebSocket):
+            await ws.accept()
+            # Wait for client message, then close with specific code
+            message = await ws.receive_text()
+            close_code = int(message)
+            await ws.close(code=close_code)
+
+    serve.run(WebSocketServer.bind())
+    url = f"{get_application_url(is_websocket=True)}/"
+
+    # Test normal close (code 1000)
+    with connect(url) as websocket:
+        websocket.send("1000")
+        with pytest.raises(ConnectionClosed):
+            websocket.recv()
+
+    # Test abnormal close (code 1011)
+    with connect(url) as websocket:
+        websocket.send("1011")
+        with pytest.raises(ConnectionClosed):
+            websocket.recv()
+
+
+def test_unary_streaming_websocket_same_deployment(serve_instance):
+    app = FastAPI()
+
+    signal_actor = SignalActor.remote()
+
+    @serve.deployment
+    @serve.ingress(app)
+    class RenaissanceMan:
+        @app.get("/")
+        def say_hi(self):
+            return "hi"
+
+        @app.get("/stream")
+        def gen_hi(self) -> StreamingResponse:
+            def gen():
+                for i in range(5):
+                    yield "hi"
+                    ray.get(signal_actor.wait.remote())
+
+            return StreamingResponse(gen(), media_type="text/plain")
+
+        @app.websocket("/ws")
+        async def ws_hi(self, ws: WebSocket):
+            try:
+                await ws.accept()
+                await ws.send_text(await ws.receive_text())
+            except WebSocketDisconnect:
+                pass
+
+    serve.run(RenaissanceMan.bind())
+
+    http_url = get_application_url()
+    assert httpx.get(http_url).json() == "hi"
+
+    with httpx.stream("GET", f"{http_url}/stream") as r:
+        r.raise_for_status()
+        for chunk in r.iter_text():
+            assert chunk == "hi"
+            ray.get(signal_actor.send.remote())
+
+    url = get_application_url(is_websocket=True)
+    with connect(f"{url}/ws") as ws:
+        ws.send("hi")
+        assert ws.recv() == "hi"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))