Building a Real-time Chat Application with WebSockets and FastAPI¶

# requirements.txt
fastapi==0.104.1
uvicorn[standard]==0.24.0
websockets==12.0
python-multipart==0.0.6
python-jose[cryptography]==3.3.0
passlib[bcrypt]==1.7.4
redis==5.0.1
pydantic==2.5.0
pytest==7.4.3
pytest-asyncio==0.21.1

# app/main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import uvicorn

from routers import auth, chat
from services.websocket_manager import ConnectionManager

app = FastAPI(title="Real-time Chat API", version="1.0.0")

# CORS middleware for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include routers
app.include_router(auth.router, prefix="/api/auth", tags=["authentication"])
app.include_router(chat.router, prefix="/api/chat", tags=["chat"])

# Serve static files
app.mount("/static", StaticFiles(directory="frontend"), name="static")

# WebSocket connection manager
manager = ConnectionManager()

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket, client_id)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.handle_message(client_id, data)
    except WebSocketDisconnect:
        manager.disconnect(client_id)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# app/services/websocket_manager.py
import json
import asyncio
from typing import Dict, List
from fastapi import WebSocket
from models.message import Message, MessageType
from services.message_service import MessageService

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.rooms: Dict[str, List[str]] = {}
        self.message_service = MessageService()

    async def connect(self, websocket: WebSocket, client_id: str):
        await websocket.accept()
        self.active_connections[client_id] = websocket
        await self.broadcast_user_joined(client_id)

    def disconnect(self, client_id: str):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
            self._remove_from_rooms(client_id)
            asyncio.create_task(self.broadcast_user_left(client_id))

    async def handle_message(self, client_id: str, message: str):
        try:
            data = json.loads(message)
            message_type = data.get("type", "message")

            if message_type == "join_room":
                await self.join_room(client_id, data["room"])
            elif message_type == "leave_room":
                await self.leave_room(client_id, data["room"])
            elif message_type == "message":
                await self.send_room_message(client_id, data["room"], data["content"])
            elif message_type == "typing":
                await self.broadcast_typing(client_id, data["room"], data["is_typing"])

        except json.JSONDecodeError:
            await self.send_error(client_id, "Invalid message format")
        except Exception as e:
            await self.send_error(client_id, f"Error processing message: {str(e)}")

    async def join_room(self, client_id: str, room: str):
        if room not in self.rooms:
            self.rooms[room] = []
        if client_id not in self.rooms[room]:
            self.rooms[room].append(client_id)
        await self.broadcast_to_room(room, {
            "type": "user_joined",
            "user": client_id,
            "room": room
        })

    async def send_room_message(self, client_id: str, room: str, content: str):
        if room not in self.rooms or client_id not in self.rooms[room]:
            await self.send_error(client_id, "Not in room")
            return

        message = Message(
            sender=client_id,
            room=room,
            content=content,
            message_type=MessageType.MESSAGE
        )

        # Save message to database
        await self.message_service.save_message(message)

        # Broadcast to room
        await self.broadcast_to_room(room, {
            "type": "message",
            "sender": client_id,
            "content": content,
            "timestamp": message.timestamp.isoformat()
        })

    async def broadcast_to_room(self, room: str, message: dict):
        if room in self.rooms:
            for client_id in self.rooms[room]:
                if client_id in self.active_connections:
                    try:
                        await self.active_connections[client_id].send_text(json.dumps(message))
                    except:
                        # Remove dead connections
                        del self.active_connections[client_id]
                        self._remove_from_rooms(client_id)

    async def send_error(self, client_id: str, error: str):
        if client_id in self.active_connections:
            await self.active_connections[client_id].send_text(json.dumps({
                "type": "error",
                "message": error
            }))

# app/models/message.py
from pydantic import BaseModel
from datetime import datetime
from enum import Enum
from typing import Optional

class MessageType(str, Enum):
    MESSAGE = "message"
    JOIN = "join"
    LEAVE = "leave"
    TYPING = "typing"

class Message(BaseModel):
    id: Optional[str] = None
    sender: str
    room: str
    content: str
    message_type: MessageType = MessageType.MESSAGE
    timestamp: datetime = datetime.utcnow()

class ChatRoom(BaseModel):
    name: str
    description: Optional[str] = None
    created_by: str
    created_at: datetime = datetime.utcnow()
    is_private: bool = False

class User(BaseModel):
    id: str
    username: str
    email: str
    is_online: bool = False
    last_seen: Optional[datetime] = None

# app/routers/auth.py
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from jose import JWTError, jwt
from passlib.context import CryptContext
from datetime import datetime, timedelta
from typing import Optional
import os

router = APIRouter()

# Security configuration
SECRET_KEY = os.getenv("SECRET_KEY", "your-secret-key-here")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

@router.post("/register")
async def register(username: str, email: str, password: str):
    # In production, use a proper database
    hashed_password = get_password_hash(password)
    # Store user in database
    return {"message": "User created successfully"}

@router.post("/login")
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
    # Verify user credentials
    # In production, check against database
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": form_data.username}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}

async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        return username
    except JWTError:
        raise credentials_exception

# tests/test_websocket.py
import pytest
import asyncio
import json
from fastapi.testclient import TestClient
from app.main import app

@pytest.fixture
def client():
    return TestClient(app)

@pytest.mark.asyncio
async def test_websocket_connection():
    with client.websocket_connect("/ws/testuser") as websocket:
        # Test connection
        assert websocket is not None

        # Test message sending
        test_message = {
            "type": "message",
            "room": "general",
            "content": "Hello, world!"
        }
        websocket.send_text(json.dumps(test_message))

        # Test message receiving
        data = websocket.receive_text()
        message = json.loads(data)
        assert message["type"] == "message"
        assert message["content"] == "Hello, world!"

@pytest.mark.asyncio
async def test_room_joining():
    with client.websocket_connect("/ws/testuser") as websocket:
        # Join room
        join_message = {
            "type": "join_room",
            "room": "test_room"
        }
        websocket.send_text(json.dumps(join_message))

        # Verify room join confirmation
        data = websocket.receive_text()
        message = json.loads(data)
        assert message["type"] == "user_joined"
        assert message["room"] == "test_room"

@pytest.mark.asyncio
async def test_typing_indicator():
    with client.websocket_connect("/ws/testuser") as websocket:
        # Send typing indicator
        typing_message = {
            "type": "typing",
            "room": "general",
            "is_typing": True
        }
        websocket.send_text(json.dumps(typing_message))

        # Verify typing indicator
        data = websocket.receive_text()
        message = json.loads(data)
        assert message["type"] == "typing"
        assert message["is_typing"] == True

# app/services/message_service.py
import redis
from typing import List
from models.message import Message
import json

class MessageService:
    def __init__(self):
        self.redis_client = redis.Redis(host='localhost', port=6379, db=0)

    async def save_message(self, message: Message):
        # Save to Redis for fast access
        key = f"room:{message.room}:messages"
        self.redis_client.lpush(key, json.dumps(message.dict()))

        # Keep only last 100 messages per room
        self.redis_client.ltrim(key, 0, 99)

    async def get_room_messages(self, room: str, limit: int = 50) -> List[Message]:
        key = f"room:{room}:messages"
        messages_data = self.redis_client.lrange(key, 0, limit - 1)
        return [Message(**json.loads(msg)) for msg in messages_data]

# app/services/presence_service.py
import redis
from datetime import datetime, timedelta
from typing import Set

class PresenceService:
    def __init__(self):
        self.redis_client = redis.Redis(host='localhost', port=6379, db=0)

    async def update_user_presence(self, user_id: str, room: str):
        key = f"presence:{room}"
        self.redis_client.sadd(key, user_id)
        self.redis_client.expire(key, 300)  # 5 minutes

    async def get_online_users(self, room: str) -> Set[str]:
        key = f"presence:{room}"
        return self.redis_client.smembers(key)

    async def remove_user_presence(self, user_id: str, room: str):
        key = f"presence:{room}"
        self.redis_client.srem(key, user_id)

# Essential WebSocket patterns
websocket_patterns = {
    "connection": "Accept connections and manage client state",
    "messaging": "Handle real-time message broadcasting",
    "rooms": "Implement room-based communication",
    "presence": "Track user online/offline status",
    "typing": "Show typing indicators",
    "error_handling": "Graceful error handling and reconnection"
}