Background Tasks & WebSockets

import logging
import httpx
from fastapi import FastAPI, BackgroundTasks, Depends
from pydantic import BaseModel, EmailStr

app = FastAPI()
logger = logging.getLogger(__name__)

# ── Task functions ──
async def send_welcome_email(email: str, username: str) -> None:
    """Async task — runs after response is sent."""
    async with httpx.AsyncClient() as client:
        await client.post("https://api.emailprovider.com/send", json={
            "to": email,
            "subject": "Welcome!",
            "body": f"Hi {username}, welcome aboard!",
        })
    logger.info("Welcome email sent to %s", email)

def write_audit_log(user_id: int, action: str) -> None:
    """Sync task — runs in threadpool so it doesn't block the loop."""
    logger.info("AUDIT user=%d action=%s", user_id, action)

# ── Route ──
class SignupBody(BaseModel):
    username: str
    email: EmailStr

@app.post("/signup", status_code=201)
async def signup(body: SignupBody, tasks: BackgroundTasks):
    user = {"id": 42, "username": body.username, "email": body.email}

    # Queue tasks — they run AFTER the response is returned to the client
    tasks.add_task(send_welcome_email, body.email, body.username)
    tasks.add_task(write_audit_log, user["id"], "signup")

    return {"user": user, "message": "Account created"}

pip install arq redis

# worker.py
import asyncio
from arq import create_pool
from arq.connections import RedisSettings

REDIS_SETTINGS = RedisSettings(host="localhost", port=6379)

# ── Task functions (must be async) ──
async def send_email(ctx, email: str, subject: str, body: str) -> dict:
    """ctx is the ARQ worker context (contains redis, job_id, etc.)"""
    print(f"Sending email to {email}: {subject}")
    await asyncio.sleep(0.5)   # simulate SMTP
    return {"sent": True, "to": email}

async def generate_report(ctx, report_id: int) -> dict:
    print(f"Generating report {report_id}...")
    await asyncio.sleep(2.0)   # simulate heavy work
    return {"report_id": report_id, "status": "done"}

# ── Worker class — defines which tasks this worker handles ──
class WorkerSettings:
    functions = [send_email, generate_report]
    redis_settings = REDIS_SETTINGS
    max_jobs = 10           # concurrent jobs
    job_timeout = 60        # seconds before a job is considered failed

# main.py — enqueue jobs from FastAPI
from fastapi import FastAPI, Request
from arq import create_pool
from arq.connections import RedisSettings
from contextlib import asynccontextmanager

REDIS_SETTINGS = RedisSettings()

@asynccontextmanager
async def lifespan(app):
    app.state.arq = await create_pool(REDIS_SETTINGS)
    yield
    await app.state.arq.aclose()

app = FastAPI(lifespan=lifespan)

@app.post("/reports")
async def request_report(report_id: int, request: Request):
    job = await request.app.state.arq.enqueue_job(
        "generate_report",   # function name as string
        report_id,
    )
    return {"job_id": job.job_id, "status": "queued"}

@app.get("/reports/{job_id}")
async def check_report(job_id: str, request: Request):
    job = await request.app.state.arq.jobs(job_id)
    return {"job_id": job_id, "status": job[0].status if job else "not_found"}

arq worker.WorkerSettings

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

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

# ── Receiving different message types ──
@app.websocket("/ws/typed")
async def typed_ws(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            # receive_text(), receive_bytes(), receive_json()
            msg = await websocket.receive_json()
            event = msg.get("event")
            if event == "ping":
                await websocket.send_json({"event": "pong"})
            elif event == "message":
                await websocket.send_json({"event": "ack", "data": msg.get("data")})
    except WebSocketDisconnect:
        pass

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import Optional
import json

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active: list[WebSocket] = []

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        self.active.append(ws)

    def disconnect(self, ws: WebSocket) -> None:
        self.active.remove(ws)

    async def send_personal(self, message: str, ws: WebSocket) -> None:
        await ws.send_text(message)

    async def broadcast(self, message: str) -> None:
        dead: list[WebSocket] = []
        for ws in self.active:
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.active.remove(ws)

    async def broadcast_json(self, data: dict) -> None:
        await self.broadcast(json.dumps(data))

manager = ConnectionManager()

@app.websocket("/ws/{client_id}")
async def chat(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    await manager.broadcast_json({"event": "join", "client": client_id})
    try:
        while True:
            text = await websocket.receive_text()
            await manager.broadcast_json({
                "event": "message",
                "client": client_id,
                "text": text,
            })
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast_json({"event": "leave", "client": client_id})

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query, status
import jwt

app = FastAPI()
SECRET_KEY = "change-me"
ALGORITHM  = "HS256"

def decode_token(token: str) -> dict:
    try:
        return jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
    except jwt.InvalidTokenError:
        return {}

@app.websocket("/ws/secure")
async def secure_ws(
    websocket: WebSocket,
    token: str = Query(...),           # ?token=<jwt>
):
    payload = decode_token(token)
    if not payload:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    user_id = payload.get("sub")
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_json({"user": user_id, "echo": data})
    except WebSocketDisconnect:
        pass

pip install redis asyncio-redis

# broadcast.py — Redis Pub/Sub connection manager
import asyncio, json, redis.asyncio as aioredis
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from contextlib import asynccontextmanager

CHANNEL = "chat"

class RedisBroadcastManager:
    def __init__(self):
        self.local: list[WebSocket] = []
        self.redis: aioredis.Redis | None = None
        self.pubsub = None

    async def startup(self):
        self.redis  = aioredis.from_url("redis://localhost")
        self.pubsub = self.redis.pubsub()
        await self.pubsub.subscribe(CHANNEL)
        asyncio.create_task(self._listener())

    async def _listener(self):
        """Forward every Redis message to all local connections."""
        async for msg in self.pubsub.listen():
            if msg["type"] == "message":
                text = msg["data"].decode()
                dead = []
                for ws in self.local:
                    try:
                        await ws.send_text(text)
                    except Exception:
                        dead.append(ws)
                for ws in dead:
                    self.local.remove(ws)

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.local.append(ws)

    def disconnect(self, ws: WebSocket):
        self.local.remove(ws)

    async def publish(self, data: dict):
        """Publish to Redis — reaches all workers."""
        await self.redis.publish(CHANNEL, json.dumps(data))

manager = RedisBroadcastManager()

@asynccontextmanager
async def lifespan(app):
    await manager.startup()
    yield

app = FastAPI(lifespan=lifespan)

@app.websocket("/ws/{client_id}")
async def chat(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    await manager.publish({"event": "join", "client": client_id})
    try:
        while True:
            text = await websocket.receive_text()
            await manager.publish({"event": "message", "client": client_id, "text": text})
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.publish({"event": "leave", "client": client_id})

from contextlib import asynccontextmanager
from fastapi import FastAPI
import asyncio, logging

logger = logging.getLogger(__name__)

async def periodic_cleanup():
    """Background coroutine that runs for the app's lifetime."""
    while True:
        logger.info("Running cleanup...")
        await asyncio.sleep(60)   # every 60 seconds

@asynccontextmanager
async def lifespan(app):
    # ── Startup ──
    logger.info("App starting up")
    task = asyncio.create_task(periodic_cleanup())
    app.state.cleanup_task = task
    yield
    # ── Shutdown ──
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass
    logger.info("App shut down cleanly")

app = FastAPI(lifespan=lifespan)

@app.get("/health")
def health():
    return {"status": "ok"}

# test_ws.py
from fastapi.testclient import TestClient
from main import app   # your app with @app.websocket("/ws/{client_id}")

client = TestClient(app)

def test_echo():
    with client.websocket_connect("/ws/test-user") as ws:
        ws.send_json({"event": "message", "text": "hello"})
        data = ws.receive_json()
        assert data["event"] == "message"
        assert data["text"] == "hello"

def test_disconnect():
    with client.websocket_connect("/ws/user-1") as ws:
        ws.send_text("ping")
        _ = ws.receive_text()
    # Connection closes cleanly when the context manager exits

def test_unauthorized():
    """Secure endpoint should reject missing/invalid token."""
    with client.websocket_connect("/ws/secure?token=bad-token") as ws:
        # Server closes the connection immediately
        import pytest
        from starlette.websockets import WebSocketDisconnect
        with pytest.raises((WebSocketDisconnect, Exception)):
            ws.receive_text()