FastAPI WebSocket: endpoint, menedżer broadcastu i testy | We.Inc
Napisz endpoint WebSocket w FastAPI, rozsyłaj wiadomości do wszystkich klientów przez connection manager, odczytuj tokeny z query stringa, obsługuj rozłączenia i testuj wszystko przez TestClient.
W FastAPI deklarujesz trasę gniazda przez @app.websocket("/ws") i asynchroniczną funkcję przyjmującą WebSocket. Wywołaj await websocket.accept(), a potem w pętli receive_text() i send_text(). Przechwyć WebSocketDisconnect, aby posprzątać, gdy klient odejdzie. Najpierw zainstaluj pakiet websockets, a testuj przez TestClient.websocket_connect().
Metody WebSocket, których będziesz używać najczęściej
Wywołanie
Co robi
await websocket.accept()
Kończy handshake; wcześniej nie można niczego wysłać
await websocket.receive_text()
Czeka na następną wiadomość tekstową
await websocket.receive_json()
Czeka na wiadomość i parsuje ją jako JSON
await websocket.send_text(s)
Wysyła wiadomość tekstową
await websocket.send_json(obj)
Serializuje obiekt i go wysyła
await websocket.close(code=1000)
Kończy połączenie z kodem zamknięcia
Poniższy kreator endpointu tworzy działający main.py na podstawie kilku wyborów: ścieżki trasy, tego, czy wiadomości to zwykły tekst czy JSON, czy każdy klient słyszy każdą wiadomość oraz czy w query stringu musi przyjść token. Dołączony jest pasujący plik pytest. Każda opcja odpowiada wzorcowi z dokumentacji FastAPI, przepisanemu tutaj jako jeden plik gotowy do uruchomienia.
Wybierz opcje w kreatorze i skopiuj main.py do pustego folderu.
Zainstaluj potrzebne pakiety: pip install "fastapi[standard]" websockets, a potem uruchom aplikację poleceniem fastapi dev main.py.
Połącz się z testera WebSocket We.Inc lub z konsoli przeglądarki pod adresem ws://localhost:8000 plus Twoja ścieżka.
Skopiuj wygenerowany plik testów i uruchom pytest, aby sprawdzić endpoint bez uruchamiania serwera.
Opakuj pętlę odbioru w try i except WebSocketDisconnect. Bez tego każdy klient, który zamknie kartę, zostawi stack trace w logach.
Zamknięcie gniazda przed accept() sprawia, że serwer odpowiada na handshake kodem HTTP 403, więc klient nigdy się nie otworzy. Najpierw zaakceptuj, a potem zamknij kodem 1008, jeśli chcesz, aby klient otrzymał czytelny kod naruszenia zasad.
Lista połączeń w pamięci zna tylko klientów bieżącego procesu. Uruchom kilka workerów, a każdy będzie rozsyłał do innego podzbioru.
Używaj receive_json i send_json, gdy wiadomości zawierają więcej niż jedno pole. Oszczędzają wywołanie json.loads i głośno zawodzą przy błędnych danych.
Klienci w przeglądarce nie mogą dodać nagłówka Authorization do gniazda, więc odczytuj dane uwierzytelniające z parametru query lub ciasteczka.
Najczęściej zadawane pytania
Dlaczego mój WebSocket w FastAPI zwraca 404 lub nie wykonuje upgrade?
Zacznij od instalacji: dokumentacja WebSocket w FastAPI każe dodać pakiet websockets obok FastAPI. Potem sprawdź ścieżkę, bo trasa zadeklarowana przez @app.websocket odpowiada tylko na upgrade do gniazda, nigdy na zwykłe żądania HTTP GET. Zamknięcie gniazda przed accept() też objawia się jako odrzucony handshake (HTTP 403).
Jak rozesłać wiadomość do wszystkich klientów?
Trzymaj zaakceptowane gniazda na liście w małej klasie menedżera, dodawaj przy połączeniu, usuwaj przy rozłączeniu i w pętli po liście wywołuj send_text. Kreator powyżej generuje dokładnie to, gdy zaznaczysz opcję broadcastu.
Czy broadcast działa z kilkoma workerami Uvicorn?
Nie z listą w pamięci. Każdy worker to osobny proces z własną listą. Dokumentacja FastAPI odsyła do encode/broadcaster, który przekazuje wiadomości przez Redis lub PostgreSQL, aby każdy proces słyszał każdą wiadomość.
Czy mogę używać Depends i Query w trasie WebSocket?
Tak. Endpointy WebSocket przyjmują Depends, Query, Cookie, Header, Path i Security tak samo jak trasy HTTP, więc zależność uwierzytelniająca może być współdzielona przez oba rodzaje.
Jak przetestować endpoint WebSocket?
Utwórz TestClient dla aplikacji i otwórz gniazdo w bloku with przez client.websocket_connect("/ws"). Wewnątrz send_text i receive_text prowadzą rozmowę, a wyjście z bloku zamyka połączenie.
Co robi pętla endpointu
Handler gniazda to długo działająca korutyna. Akceptuje połączenie raz, a potem czeka na await receive_text() do następnej wiadomości, więc pojedynczy klient w bezczynności prawie nic nie kosztuje. Gdy klient odejdzie, oczekujące receive rzuca WebSocketDisconnect i to jest moment na sprzątanie.
Ponieważ handler jest asynchroniczny, inni klienci są obsługiwani, gdy jeden czeka. Unikaj w pętli blokujących wywołań, takich jak time.sleep czy synchroniczne sterowniki baz danych; zamrażają one każde połączenie w danym workerze.
Uwierzytelnianie przed accept
Odczytaj token z query stringa jako zwykły parametr. Masz wtedy dwa sposoby, by go odrzucić. Zamknięcie przed accept() zamienia się w HTTP 403 przy handshake, które przeglądarka zgłasza tylko jako nieudane połączenie. Zaakceptowanie i zamknięcie z kodem 1008 pozwala klientowi odczytać powód i pokazać prośbę o zalogowanie zamiast ponawiać próby w nieskończoność.
Najmniejszy działający endpoint
main.py: echo
Zaakceptuj, a potem odsyłaj każdą wiadomość, dopóki klient się nie rozłączy.
Broadcast z connection managerem
Klasa menedżera współdzielona przez wszystkie połączenia
Działa w jednym procesie. Zanim zaczniesz skalować, zamień listę na przekaźnik oparty na Redis.
Testowanie przez TestClient
Serwer nie jest potrzebny: klient rozmawia z aplikacją w tym samym procesie. Uruchom go przez pytest.
Kreator endpointu WebSocket dla FastAPI
Wybierz, jak ma działać endpoint, i skopiuj kompletny main.py oraz plik pytest, który go sprawdza.
Format wiadomości
Zwykły tekst
JSON
Wysyłaj każdą wiadomość do każdego połączonego klienta
Wymagaj tokenu w query stringu (?token=...)
Uruchom: fastapi dev main.py, a potem połącz się z ws://localhost:8000{v}