FastAPI WebSocket Testing

Introduction

Testing is a crucial part of developing robust applications, and when it comes to WebSockets in FastAPI, proper testing ensures that your real-time communication features work reliably. In this guide, we'll explore different approaches to testing WebSocket connections in FastAPI, from unit tests to integration tests, using tools like Pytest and the TestClient.

WebSocket testing differs from standard HTTP endpoint testing because:

WebSockets maintain persistent connections
They involve bidirectional communication
They often include connection state management
Testing typically requires simulating both server and client behaviors

By the end of this guide, you'll understand how to effectively test your FastAPI WebSocket implementations to ensure they work correctly before deployment.

Prerequisites

Before diving in, ensure you have:

Basic knowledge of FastAPI
Understanding of WebSockets in FastAPI
Python 3.7+ installed
FastAPI and its dependencies installed

Let's install the necessary packages:

bash
pip install fastapi uvicorn websockets pytest pytest-asyncio httpx

Setting Up a Basic WebSocket Application for Testing

First, let's create a simple FastAPI application with a WebSocket endpoint that we'll use for testing:

python
# app.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            # Echo the received message with a prefix
            await websocket.send_text(f"Message received: {data}")
    except WebSocketDisconnect:
        print("Client disconnected")

This simple WebSocket endpoint accepts connections and echoes back any messages it receives with a prefix.

Testing Approaches for WebSockets

1. Unit Testing with TestClient

FastAPI's TestClient doesn't directly support WebSocket testing, but we can extend it or use alternative approaches. Let's see how we can test our WebSocket endpoint using the websockets library:

python
# test_websocket.py
import asyncio
import pytest
from fastapi.testclient import TestClient
import websockets
import uvicorn
import threading
import time
from app import app

# Start the server in a separate thread for testing
@pytest.fixture
def server():
    server_thread = threading.Thread(
        target=uvicorn.run,
        args=(app,),
        kwargs={"host": "127.0.0.1", "port": 8000, "log_level": "critical"},
        daemon=True,
    )
    server_thread.start()
    time.sleep(1)  # Give the server time to start
    yield
    # No need to stop - daemon thread will be terminated when the test ends

@pytest.mark.asyncio
async def test_websocket_echo(server):
    # Connect to the WebSocket endpoint
    uri = "ws://127.0.0.1:8000/ws"
    async with websockets.connect(uri) as websocket:
        # Send a message
        test_message = "Hello WebSocket"
        await websocket.send(test_message)
        
        # Receive the response
        response = await websocket.recv()
        
        # Assert that the response is as expected
        assert response == f"Message received: {test_message}"

2. Using ASGI TestClient with httpx

A more integrated approach is to use the httpx library which provides an ASGI test client capable of handling WebSocket connections:

python
# test_websocket_httpx.py
import pytest
import asyncio
from httpx import AsyncClient
import json
from app import app

@pytest.mark.asyncio
async def test_websocket_with_httpx():
    async with AsyncClient(app=app, base_url="http://test") as client:
        # Note: httpx doesn't have native WebSocket support, so for real WebSocket testing
        # we'd need additional tools. This example shows API testing pattern.
        response = await client.get("/")
        assert response.status_code == 404  # We didn't define a root route

Since httpx doesn't provide native WebSocket testing, we would typically combine it with a dedicated WebSocket client or library.

3. Using a Custom WebSocket Test Client

For more control over WebSocket testing, we can create a custom test client:

python
# test_websocket_custom.py
import pytest
from fastapi import FastAPI, WebSocket
from fastapi.testclient import TestClient
import asyncio

class WebSocketTestClient:
    def __init__(self, app: FastAPI, path: str):
        self.app = app
        self.path = path
        self.client = TestClient(app)
        self.incoming_messages = asyncio.Queue()
        self.outgoing_messages = asyncio.Queue()
        
    async def connect(self):
        # This would be a stub for simulation purposes
        pass
        
    async def send_text(self, data: str):
        # In a real implementation, this would send data to the server
        await self.outgoing_messages.put(data)
        
    async def receive_text(self):
        # In a real implementation, this would wait for server messages
        return await self.incoming_messages.get()
        
    async def close(self):
        # Close the connection
        pass

# Example test using our custom client (simplified)
@pytest.mark.asyncio
async def test_echo():
    from app import app
    
    client = WebSocketTestClient(app, "/ws")
    await client.connect()
    
    # Simulate sending message
    test_message = "Test message"
    await client.send_text(test_message)
    
    # Simulate receiving message (in a real test, we'd integrate with the actual endpoint)
    # For demonstration, we'd manually put the expected response in the queue
    await client.incoming_messages.put(f"Message received: {test_message}")
    
    # Assert the response
    response = await client.receive_text()
    assert response == f"Message received: {test_message}"
    
    await client.close()

This is a simplified example of a custom WebSocket test client. In a real-world scenario, you'd need to hook this up to actually communicate with your FastAPI WebSocket endpoints.

Creating Integration Tests with pytest-asyncio

For more comprehensive testing, we can create integration tests that simulate real client-server interactions:

python
# test_websocket_integration.py
import pytest
import asyncio
import websockets
from app import app
import uvicorn
import threading
import time

@pytest.fixture(scope="module")
def run_server():
    # Start server in a separate thread
    server_thread = threading.Thread(
        target=uvicorn.run,
        args=(app,),
        kwargs={
            "host": "127.0.0.1",
            "port": 8000,
            "log_level": "critical"
        },
        daemon=True
    )
    server_thread.start()
    time.sleep(1)  # Give server time to start
    yield
    # Thread will be terminated when test ends (daemon=True)

@pytest.mark.asyncio
async def test_websocket_communication(run_server):
    uri = "ws://127.0.0.1:8000/ws"
    
    # Connect to the WebSocket
    async with websockets.connect(uri) as websocket:
        # Test sending multiple messages
        for i in range(3):
            message = f"Test message {i}"
            await websocket.send(message)
            response = await websocket.recv()
            assert response == f"Message received: {message}"
            
        # Test receiving multiple messages in sequence
        messages = ["Hello", "World", "WebSocket"]
        for message in messages:
            await websocket.send(message)
            response = await websocket.recv()
            assert response == f"Message received: {message}"

Testing WebSocket Authentication

Many real-world applications require authentication for WebSocket connections. Let's see how to test authenticated WebSockets:

python
# auth_websocket_app.py
from fastapi import FastAPI, WebSocket, Depends, WebSocketDisconnect, HTTPException, status
from fastapi.security import APIKeyHeader

app = FastAPI()

# Simple API key authentication
API_KEY = "test_api_key"
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)

async def get_api_key(api_key: str = Depends(api_key_header)):
    if api_key == API_KEY:
        return api_key
    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Invalid API Key",
    )

@app.websocket("/secure-ws")
async def secure_websocket_endpoint(websocket: WebSocket):
    # Extract headers for authentication
    headers = dict(websocket.headers)
    api_key = headers.get("authorization")
    
    if api_key != API_KEY:
        await websocket.close(code=1008)  # Policy violation
        return
    
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Secure message received: {data}")
    except WebSocketDisconnect:
        print("Client disconnected from secure websocket")

Now let's test this authenticated WebSocket endpoint:

python
# test_auth_websocket.py
import pytest
import asyncio
import websockets
from auth_websocket_app import app
import uvicorn
import threading
import time

@pytest.fixture(scope="module")
def run_auth_server():
    server_thread = threading.Thread(
        target=uvicorn.run,
        args=(app,),
        kwargs={
            "host": "127.0.0.1",
            "port": 8001,
            "log_level": "critical"
        },
        daemon=True
    )
    server_thread.start()
    time.sleep(1)  # Give server time to start
    yield

@pytest.mark.asyncio
async def test_secure_websocket_with_valid_auth(run_auth_server):
    uri = "ws://127.0.0.1:8001/secure-ws"
    
    # Connect with valid API key
    headers = {
        "Authorization": "test_api_key"
    }
    
    async with websockets.connect(uri, extra_headers=headers) as websocket:
        message = "Authenticated message"
        await websocket.send(message)
        response = await websocket.recv()
        assert response == f"Secure message received: {message}"

@pytest.mark.asyncio
async def test_secure_websocket_with_invalid_auth(run_auth_server):
    uri = "ws://127.0.0.1:8001/secure-ws"
    
    # Connect with invalid API key
    headers = {
        "Authorization": "invalid_key"
    }
    
    # The connection should be closed by the server
    with pytest.raises(websockets.exceptions.ConnectionClosedError):
        async with websockets.connect(uri, extra_headers=headers) as websocket:
            await websocket.send("This should not work")

Testing WebSocket Broadcast Functionality

Let's test a more complex scenario where messages are broadcast to multiple connected clients:

python
# broadcast_app.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List

app = FastAPI()

# Store for connected clients
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
    
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/broadcast-ws")
async def broadcast_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            # Broadcast the message to all connected clients
            await manager.broadcast(f"Broadcast: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)

Now, let's test this broadcast functionality:

python
# test_broadcast.py
import pytest
import asyncio
import websockets
from broadcast_app import app
import uvicorn
import threading
import time

@pytest.fixture(scope="module")
def run_broadcast_server():
    server_thread = threading.Thread(
        target=uvicorn.run,
        args=(app,),
        kwargs={
            "host": "127.0.0.1",
            "port": 8002,
            "log_level": "critical"
        },
        daemon=True
    )
    server_thread.start()
    time.sleep(1)  # Give server time to start
    yield

@pytest.mark.asyncio
async def test_broadcast_functionality(run_broadcast_server):
    uri = "ws://127.0.0.1:8002/broadcast-ws"
    
    # Create three client connections
    async def client_connection(client_id, message_queue):
        async with websockets.connect(uri) as websocket:
            # Add this client to ready clients
            message_queue.put_nowait(f"Client {client_id} ready")
            
            # Wait for broadcast message
            response = await websocket.recv()
            message_queue.put_nowait(f"Client {client_id} received: {response}")
    
    message_queue = asyncio.Queue()
    
    # Start three client tasks
    client_tasks = [
        asyncio.create_task(client_connection(i, message_queue))
        for i in range(3)
    ]
    
    # Wait for all clients to be ready
    ready_count = 0
    while ready_count < 3:
        message = await message_queue.get()
        if "ready" in message:
            ready_count += 1
    
    # Send a broadcast message from a new client
    async with websockets.connect(uri) as broadcast_client:
        await broadcast_client.send("Hello everyone!")
        
        # This client will also receive the broadcast
        broadcast_response = await broadcast_client.recv()
        assert broadcast_response == "Broadcast: Hello everyone!"
    
    # Verify all clients received the broadcast
    received_count = 0
    while not message_queue.empty():
        message = await message_queue.get()
        if "received" in message:
            assert "Broadcast: Hello everyone!" in message
            received_count += 1
    
    # Ensure all three clients received the message
    assert received_count == 3
    
    # Clean up tasks
    for task in client_tasks:
        task.cancel()

Best Practices for WebSocket Testing

Isolate Tests: Each test should focus on a specific aspect of WebSocket functionality
Use Proper Fixtures: Set up and tear down server instances efficiently
Test Edge Cases:
- Connection handling
- Authentication failures
- Disconnections
- Message size limits
- Concurrent connections
Test Performance: For applications expecting many concurrent connections, test with a larger number of simulated clients
Mock External Dependencies: If your WebSocket handlers interact with databases or other services, consider mocking these dependencies
Test Reconnection Logic: Ensure clients can reconnect properly after disconnections

Creating a Complete Test Suite

Here's how you might structure a complete test suite for your WebSocket application:

python
# test_suite.py
import pytest
import asyncio
import websockets
from fastapi import FastAPI, WebSocket
import uvicorn
import threading
import time
from app import app

@pytest.fixture(scope="module")
def server():
    server_thread = threading.Thread(
        target=uvicorn.run,
        args=(app,),
        kwargs={"host": "127.0.0.1", "port": 8000, "log_level": "critical"},
        daemon=True,
    )
    server_thread.start()
    time.sleep(1)  # Give the server time to start
    yield

class TestWebSocketFunctionality:
    @pytest.mark.asyncio
    async def test_connection(self, server):
        """Test that a client can connect to the WebSocket endpoint."""
        uri = "ws://127.0.0.1:8000/ws"
        async with websockets.connect(uri) as websocket:
            # Connection successful if no exception
            pass
    
    @pytest.mark.asyncio
    async def test_echo_message(self, server):
        """Test that the server echoes back messages."""
        uri = "ws://127.0.0.1:8000/ws"
        async with websockets.connect(uri) as websocket:
            message = "Test echo"
            await websocket.send(message)
            response = await websocket.recv()
            assert response == f"Message received: {message}"
    
    @pytest.mark.asyncio
    async def test_multiple_messages(self, server):
        """Test sending multiple messages in sequence."""
        uri = "ws://127.0.0.1:8000/ws"
        async with websockets.connect(uri) as websocket:
            messages = ["First", "Second", "Third"]
            for msg in messages:
                await websocket.send(msg)
                response = await websocket.recv()
                assert response == f"Message received: {msg}"
    
    @pytest.mark.asyncio
    async def test_connection_close(self, server):
        """Test that the connection can be closed properly."""
        uri = "ws://127.0.0.1:8000/ws"
        websocket = await websockets.connect(uri)
        await websocket.close()
        # Check that it's closed
        assert websocket.closed

Summary

Testing WebSocket connections in FastAPI requires a different approach than testing standard HTTP endpoints. In this guide, we've covered:

Setting up basic WebSocket endpoints for testing
Using various testing approaches including:
- Custom WebSocket test clients
- Integration testing with pytest-asyncio
- Testing authenticated WebSockets
- Testing broadcast functionality
Best practices for WebSocket testing
Creating a complete test suite

By implementing proper tests for your WebSocket endpoints, you can ensure they function correctly, handle edge cases appropriately, and provide a reliable experience for your users.

Additional Resources

Exercises

Create a FastAPI WebSocket chat application and write tests to verify that messages are correctly delivered to all connected clients.
Implement rate limiting for WebSocket messages and write tests to verify that clients exceeding the rate limit are handled appropriately.
Create a WebSocket endpoint that streams real-time data and write tests to verify the streaming functionality works correctly.
Implement and test a reconnection mechanism that allows clients to resume their session after a disconnection.
Write a load test for your WebSocket application to determine how many concurrent connections it can handle.

If you spot any mistakes on this website, please let me know at [email protected]. I’d greatly appreciate your feedback! :)

Introduction​

Prerequisites​

Setting Up a Basic WebSocket Application for Testing​

Testing Approaches for WebSockets​

1. Unit Testing with TestClient​

2. Using ASGI TestClient with httpx​

3. Using a Custom WebSocket Test Client​

Creating Integration Tests with pytest-asyncio​

Testing WebSocket Authentication​

Testing WebSocket Broadcast Functionality​

Best Practices for WebSocket Testing​

Creating a Complete Test Suite​

Summary​

Additional Resources​

Exercises​