Websocket Routes Only Work on FastAPI, not APIRouter when using prefix

[brunopcarv]

• I added a very descriptive title to this issue.
• I used the GitHub search to find a similar issue and didn't find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google "How to X in FastAPI" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.
• After submitting this, I commit to one of:
  • Read open issues with questions until I find 2 issues where I can help someone and add a comment to help there.
  • I already hit the "watch" button in this repository to receive notifications and I commit to help at least 2 people that ask questions in the future.
  • Implement a Pull Request for a confirmed bug.

Description

Websocket routes appear to only work on the main FastAPI object, not on APIRouter objects (instead it raises WebSocketDisconnect).

Example

Here's a self-contained, minimal, reproducible, example with my use case:

from fastapi import FastAPI
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocket
from fastapi import APIRouter 

app = FastAPI()
router = APIRouter(prefix="/test")
app.include_router(router)


@app.websocket_route("/ws")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"msg": "Hello WebSocket"})
    await websocket.close()


@router.websocket_route("/ws")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"msg": "Hello WebSocket 2"})
    await websocket.close()


def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket"}


def test_websocket2():
    client = TestClient(app)
    with client.websocket_connect("/test/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket 2"}

Further info

• Only the second websocket fails in the example above.
• Error log:

def test_websocket2():
        client = TestClient(app)
>      with client.websocket_connect("/test/ws") as websocket:

/opt/conda/envs/build_env/envs/run_env/lib/python3.6/site-packages/starlette/testclient.py:444: in websocket_connect
super().request("GET", url, **kwargs)
/opt/conda/envs/build_env/envs/run_env/lib/python3.6/site-packages/requests/sessions.py:542: in request
resp = self.send(prep, **send_kwargs)
/opt/conda/envs/build_env/envs/run_env/lib/python3.6/site-packages/requests/sessions.py:655: in send
r = adapter.send(request, **kwargs)
/opt/conda/envs/build_env/envs/run_env/lib/python3.6/site-packages/starlette/testclient.py:145: in send
session = WebSocketTestSession(self.app, scope)
/opt/conda/envs/build_env/envs/run_env/lib/python3.6/site-packages/starlette/testclient.py:278: in __init__
self._raise_on_close(message)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
self = <starlette.testclient.WebSocketTestSession object at 0x7f7cbc689c50>
message = {'code': 1000, 'type': 'websocket.close'}
def _raise_on_close(self, message: Message) -> None:
        if message["type"] == "websocket.close":
>           raise WebSocketDisconnect(message.get("code", 1000))
E           starlette.websockets.WebSocketDisconnect: 1000

Environment

• OS: Amazon Linux Workspace
• FastAPI Version: 0.63.0
• Python version: 3.6.8

[Kludex]

Hey @brunopcarv !

Yes, you're completely right. It's a bug introduced on the current FastAPI version (0.63.0). But it's already being addressed on this PR: #2640 .

For now, your code already works if you use this implementation instead:

from fastapi import FastAPI
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocket
from fastapi import APIRouter 

app = FastAPI()
router = APIRouter()


@app.websocket_route("/ws")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"msg": "Hello WebSocket"})
    await websocket.close()


@router.websocket_route("/ws")
async def websocket(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"msg": "Hello WebSocket 2"})
    await websocket.close()


app.include_router(router, prefix="/test")


def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket"}


def test_websocket2():
    client = TestClient(app)
    with client.websocket_connect("/test/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket 2"}

Please notice that the router is included after you define the routes. 😉

As this is duplicated, do you mind closing the issue (#2639)? 😗

[brunopcarv]

@Kludex thank you very much for your help!

[tiangolo]

Thanks for the help here @Kludex ! 👏 🙇

Thanks for reporting back and closing the issue @brunopcarv 👍

Sorry for the long delay! 🙈 I wanted to personally address each issue/PR and they piled up through time, but now I'm checking each one in order.