From bed8cfb78bdfb1502584536d12d80f965700b4c7 Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Date: Mon, 3 Feb 2025 17:50:34 -0300 Subject: [PATCH 1/4] Fixes WebSocket Implementation for Router of Instance of APIWebSocketRoute --- fastapi_utils/cbv.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/fastapi_utils/cbv.py b/fastapi_utils/cbv.py index 231325c5..c24bd416 100644 --- a/fastapi_utils/cbv.py +++ b/fastapi_utils/cbv.py @@ -13,7 +13,7 @@ import pydantic from fastapi import APIRouter, Depends -from fastapi.routing import APIRoute +from fastapi.routing import APIRoute, APIWebSocketRoute from starlette.routing import Route, WebSocketRoute PYDANTIC_VERSION = pydantic.VERSION @@ -108,12 +108,16 @@ def _register_endpoints(router: APIRouter, cls: Type[Any], *urls: str) -> None: _allocate_routes_by_method_name(router, url, function_members) router_roles = [] for route in router.routes: - if not isinstance(route, APIRoute): - raise ValueError("The provided routes should be of type APIRoute") + if not isinstance(route, APIRoute) and not isinstance(route, APIWebSocketRoute): + raise ValueError("The provided routes should be of type APIRoute or APIWebSocketRoute") - route_methods: Any = route.methods - cast(Tuple[Any], route_methods) - router_roles.append((route.path, tuple(route_methods))) + if isinstance(route, APIRoute): + route_methods: Any = route.methods + cast(Tuple[Any], route_methods) + router_roles.append((route.path, tuple(route_methods))) + + if isinstance(route, APIWebSocketRoute): + router_roles.append((route.path, tuple(["WS"]))) if len(set(router_roles)) != len(router_roles): raise Exception("An identical route role has been implemented more then once") From a525f9d455f063a8e53046a048db43cb214e4fa9 Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Date: Mon, 24 Feb 2025 18:42:40 -0300 Subject: [PATCH 2/4] Adds TestCase for FastAPIUtils CBV and Websockets --- tests/test_cbv.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/test_cbv.py b/tests/test_cbv.py index b5b9f646..8c083a37 100644 --- a/tests/test_cbv.py +++ b/tests/test_cbv.py @@ -3,8 +3,9 @@ from typing import Any, ClassVar, Optional import pytest -from fastapi import APIRouter, Depends, Request +from fastapi import APIRouter, Depends, Request, WebSocket from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect from fastapi_utils.cbv import cbv @@ -147,3 +148,19 @@ def example(self, request: Request) -> str: client = TestClient(router) response = client.get("/foo") assert response.json() == "http://testserver/bar" + + def test_websocket_router(self, router: APIRouter) -> None: + @cbv(router) + class Foo: + @router.websocket("/ws") + async def example(self, websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_text("hello") + await websocket.close() + + client = TestClient(router) + with client.websocket_connect("/ws") as websocket: + assert websocket.receive_text() == "hello" + + with pytest.raises(WebSocketDisconnect): + assert websocket.receive_text() From a224089264b3d328861e02f2e474e0f6eb462ec2 Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Date: Mon, 24 Feb 2025 18:27:04 -0300 Subject: [PATCH 3/4] Alters Router Instance Check for Better Readability --- fastapi_utils/cbv.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_utils/cbv.py b/fastapi_utils/cbv.py index c24bd416..677b9316 100644 --- a/fastapi_utils/cbv.py +++ b/fastapi_utils/cbv.py @@ -108,7 +108,7 @@ def _register_endpoints(router: APIRouter, cls: Type[Any], *urls: str) -> None: _allocate_routes_by_method_name(router, url, function_members) router_roles = [] for route in router.routes: - if not isinstance(route, APIRoute) and not isinstance(route, APIWebSocketRoute): + if not (isinstance(route, APIRoute) or isinstance(route, APIWebSocketRoute)): raise ValueError("The provided routes should be of type APIRoute or APIWebSocketRoute") if isinstance(route, APIRoute): From 5e2b161f3e9515cce77d7afecdf5d678e409c5b7 Mon Sep 17 00:00:00 2001 From: Vitor Hideyoshi Date: Sat, 8 Mar 2025 00:40:42 -0300 Subject: [PATCH 4/4] Adds Simple Receive Test --- tests/test_cbv.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_cbv.py b/tests/test_cbv.py index 8c083a37..c2ca2643 100644 --- a/tests/test_cbv.py +++ b/tests/test_cbv.py @@ -156,11 +156,13 @@ class Foo: async def example(self, websocket: WebSocket) -> None: await websocket.accept() await websocket.send_text("hello") + "VALID_VALUE" == await websocket.receive_text() await websocket.close() client = TestClient(router) with client.websocket_connect("/ws") as websocket: assert websocket.receive_text() == "hello" + websocket.send_text("VALID_VALUE") with pytest.raises(WebSocketDisconnect): assert websocket.receive_text()