wehub-resource-sync
159 lines
4.4 KiB
Python
159 lines
4.4 KiB
Python
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__]))
|