Phase 5: Redis pub/sub for horizontal scaling

Splits the WebSocket layer into three pieces so one app instance and many
behave identically: ConnectionManager stays a purely local socket registry;
RoomBroadcaster publishes chat messages to a per-room Redis channel and
every instance (including the publisher) forwards received messages to its
own local sockets via a single psubscribe("room:*") listener started in
main.py's lifespan; Presence is a Redis-backed refcounted hash per room
tracking who's connected across all instances.

Presence replaces the old process-local connected_user_ids check that
Phase 4's offline-push logic used -- without it, a user connected on a
different instance would look offline and get a redundant push. Fixing
this was scoped in beyond the issue's literal ask (message fan-out only)
since it's a real correctness gap in a phase specifically about running
more than one instance; a known limitation (no heartbeat/TTL, so a hard
crash leaks a presence increment) is documented in the README instead of
solved here.

New tests/test_broadcast.py spins up two independent app instances sharing
one Postgres + Redis to prove delivery and presence both actually cross
the Redis boundary, not just work in-process. Manually verified the same
thing against two real uvicorn processes on different ports.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-14 07:13:06 -06:00
co-authored by Claude Sonnet 5
parent d09bf4a30a
commit 0b995ef75f
11 changed files with 328 additions and 57 deletions
+1
View File
@@ -1,6 +1,7 @@
DATABASE_URL=postgresql+asyncpg://chatapp:chatapp@localhost:5432/chatapp DATABASE_URL=postgresql+asyncpg://chatapp:chatapp@localhost:5432/chatapp
SESSION_SECRET=change-me-to-a-long-random-string SESSION_SECRET=change-me-to-a-long-random-string
SESSION_HTTPS_ONLY=false SESSION_HTTPS_ONLY=false
REDIS_URL=redis://localhost:6379/0
# Optional: push notifications are skipped if unset. Generate with: # Optional: push notifications are skipped if unset. Generate with:
# python -m app.cli generate-vapid-keys # python -m app.cli generate-vapid-keys
+58 -17
View File
@@ -1,10 +1,10 @@
# KeepItTalking backend (Phase 1 + 2 + 4) # KeepItTalking backend (Phase 1 + 2 + 4 + 5)
FastAPI + SQLAlchemy 2.0 (async) + PostgreSQL. Implements auth, room CRUD FastAPI + SQLAlchemy 2.0 (async) + PostgreSQL + Redis. Implements auth, room
(open and private), room roles (owner/admin/member) and invites, a CRUD (open and private), room roles (owner/admin/member) and invites, a
single-instance WebSocket chat endpoint, and Web Push notifications for WebSocket chat endpoint that fans out across multiple app-server instances
offline room members. See `../ARCHITECTURE.md` for the full system design via Redis pub/sub, and Web Push notifications for offline room members. See
and the phased build plan. `../ARCHITECTURE.md` for the full system design and the phased build plan.
This is an **invite-only site**: there is no public registration endpoint. This is an **invite-only site**: there is no public registration endpoint.
Accounts are created by an operator on the app server — see step 4 below. Accounts are created by an operator on the app server — see step 4 below.
@@ -30,7 +30,16 @@ docker exec chatapp-postgres psql -U chatapp -d chatapp -c "CREATE DATABASE chat
(Docker here is purely a local-dev convenience for standing up Postgres quickly — (Docker here is purely a local-dev convenience for standing up Postgres quickly —
the actual deployment target has no containers at all, see `ARCHITECTURE.md` §9.) the actual deployment target has no containers at all, see `ARCHITECTURE.md` §9.)
### 2. Python environment ### 2. Redis
Used for cross-instance WebSocket fan-out and presence (see the section
below). Required — there's no in-memory fallback.
```bash
docker run -d --name chatapp-redis -p 6379:6379 redis:7-alpine
```
### 3. Python environment
```bash ```bash
cd backend cd backend
@@ -41,13 +50,13 @@ cp .env.example .env
# python3 -c "import secrets; print(secrets.token_urlsafe(32))" # python3 -c "import secrets; print(secrets.token_urlsafe(32))"
``` ```
### 3. Migrations ### 4. Migrations
```bash ```bash
.venv/bin/alembic upgrade head .venv/bin/alembic upgrade head
``` ```
### 4. Create a user ### 5. Create a user
There's no public sign-up. Create accounts directly with the CLI (add There's no public sign-up. Create accounts directly with the CLI (add
`--admin` to grant `is_site_admin`, useful ahead of the phase-6 admin portal): `--admin` to grant `is_site_admin`, useful ahead of the phase-6 admin portal):
@@ -56,7 +65,7 @@ There's no public sign-up. Create accounts directly with the CLI (add
.venv/bin/python -m app.cli create-user alice alice@example.com "some-password" .venv/bin/python -m app.cli create-user alice alice@example.com "some-password"
``` ```
### 5. (Optional) Set up push notifications ### 6. (Optional) Set up push notifications
Push works without any setup — `VAPID_PUBLIC_KEY`/`VAPID_PRIVATE_KEY` are Push works without any setup — `VAPID_PUBLIC_KEY`/`VAPID_PRIVATE_KEY` are
unset by default and push delivery is silently skipped. To enable it: unset by default and push delivery is silently skipped. To enable it:
@@ -66,7 +75,7 @@ unset by default and push delivery is silently skipped. To enable it:
# paste the three printed lines into backend/.env # paste the three printed lines into backend/.env
``` ```
### 6. Run the dev server ### 7. Run the dev server
```bash ```bash
.venv/bin/uvicorn app.main:app --reload .venv/bin/uvicorn app.main:app --reload
@@ -74,10 +83,16 @@ unset by default and push delivery is silently skipped. To enable it:
API docs: http://localhost:8000/docs. WebSocket chat endpoint: `ws://localhost:8000/ws/chat`. API docs: http://localhost:8000/docs. WebSocket chat endpoint: `ws://localhost:8000/ws/chat`.
### 7. Run tests To try horizontal scaling locally, run a second instance on another port
against the same Postgres + Redis (`.venv/bin/uvicorn app.main:app --port 8001`)
— a message sent through one instance's WebSocket is delivered to clients
connected to the other, purely via Redis.
### 8. Run tests
Tests run against a real Postgres database (`chatapp_test` by default — native Tests run against a real Postgres database (`chatapp_test` by default — native
`ENUM`/`UUID` types aren't faithfully reproduced by SQLite), with each test `ENUM`/`UUID` types aren't faithfully reproduced by SQLite) and a real Redis
(db 15 by default, kept separate from dev use of db 0), with each test
wrapped in a transaction that's rolled back afterward: wrapped in a transaction that's rolled back afterward:
```bash ```bash
@@ -99,20 +114,46 @@ app/
schemas/ Pydantic request/response models schemas/ Pydantic request/response models
routers/ auth, rooms, invites, push, health routers/ auth, rooms, invites, push, health
services/ business logic called by routers services/ business logic called by routers
ws/ WebSocket connection manager + /ws/chat handler ws/ connection_manager (local sockets), presence +
broadcaster (Redis), /ws/chat handler
alembic/ migrations alembic/ migrations
tests/ pytest + httpx/TestClient tests tests/ pytest + httpx/TestClient tests
``` ```
## Cross-instance broadcast (Phase 5)
The WebSocket layer is split into three pieces so that running one app
instance and running many behave identically:
- `app/ws/connection_manager.py` — purely local: which sockets on *this*
process are in which room, used only to actually `send_json` to them.
- `app/ws/broadcaster.py` (`RoomBroadcaster`) — on a chat message,
`publish()`s it to a Redis channel scoped to the room (`room:{id}`).
Every app instance, including the publisher, runs a single background
`listen()` task (started in `app/main.py`'s lifespan) pattern-subscribed
to `room:*`; each message it receives is handed to its own local
`ConnectionManager.broadcast()`. One instance just talks to itself
through Redis, so there's no separate single-instance code path.
- `app/ws/presence.py` (`Presence`) — a Redis hash per room
(`presence:{room_id}`, field = user ID, value = connection refcount) is
the cross-instance answer to "is this member connected *anywhere* right
now," which is what the Phase 4 offline-push check uses instead of the
local `ConnectionManager`. Refcounted so a user connected from two tabs
(or two instances) isn't marked offline until every connection closes.
Known limitation: `Presence` has no heartbeat/TTL, so a hard process crash
(not a clean disconnect) leaks that connection's increment forever — same
category of simplification as the "no server-side session revocation" note
below.
## Push notifications (Phase 4) ## Push notifications (Phase 4)
`POST /api/push/subscribe` (upserts by `endpoint`) / `DELETE /api/push/subscribe` `POST /api/push/subscribe` (upserts by `endpoint`) / `DELETE /api/push/subscribe`
manage a user's `push_subscriptions` rows; `GET /api/push/vapid-public-key` gives manage a user's `push_subscriptions` rows; `GET /api/push/vapid-public-key` gives
the frontend the key it needs for `PushManager.subscribe()`. On every chat the frontend the key it needs for `PushManager.subscribe()`. On every chat
message, `app/ws/chat.py` computes `room members - ConnectionManager. message, `app/ws/chat.py` computes `room members - Presence.
connected_user_ids(room_id)` (who's actually connected to *that room* right connected_user_ids(room_id)` (who's actually connected to *that room* right
now, tracked alongside the existing WebSocket registry) and sends each now, across every app instance — see Phase 5 below) and sends each offline
offline member a push via `pywebpush`, awaited inline against the same member a push via `pywebpush`, awaited inline against the same
request-scoped session rather than fired as a background task — the request-scoped session rather than fired as a background task — the
broadcast to online members already happened by that point, so nothing broadcast to online members already happened by that point, so nothing
online-facing is delayed, and it sidesteps `asyncio.create_task()`s outliving online-facing is delayed, and it sidesteps `asyncio.create_task()`s outliving
+4
View File
@@ -16,5 +16,9 @@ class Settings(BaseSettings):
vapid_private_key: str | None = None vapid_private_key: str | None = None
vapid_subject: str = "mailto:admin@example.com" vapid_subject: str = "mailto:admin@example.com"
# Cross-instance WebSocket fan-out + presence (ARCHITECTURE.md phase 5).
# No credentials in a local-dev default, unlike database_url.
redis_url: str = "redis://localhost:6379/0"
settings = Settings() settings = Settings()
+25 -1
View File
@@ -1,14 +1,38 @@
import asyncio
import contextlib
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastapi import FastAPI from fastapi import FastAPI
from redis.asyncio import Redis
from starlette.middleware.sessions import SessionMiddleware from starlette.middleware.sessions import SessionMiddleware
from app.config import settings from app.config import settings
from app.routers import auth, health, invites, push, rooms from app.routers import auth, health, invites, push, rooms
from app.ws.broadcaster import RoomBroadcaster
from app.ws.chat import router as ws_router from app.ws.chat import router as ws_router
from app.ws.connection_manager import ConnectionManager from app.ws.connection_manager import ConnectionManager
from app.ws.presence import Presence
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
redis = Redis.from_url(settings.redis_url, decode_responses=True)
app.state.presence = Presence(redis)
broadcaster = RoomBroadcaster(redis, app.state.connection_manager)
app.state.broadcaster = broadcaster
listener_task = asyncio.create_task(broadcaster.listen())
yield
listener_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await listener_task
await redis.aclose()
def create_app() -> FastAPI: def create_app() -> FastAPI:
app = FastAPI(title="KeepItTalking") app = FastAPI(title="KeepItTalking", lifespan=lifespan)
app.add_middleware( app.add_middleware(
SessionMiddleware, SessionMiddleware,
+41
View File
@@ -0,0 +1,41 @@
import json
import uuid
from redis.asyncio import Redis
from app.ws.connection_manager import ConnectionManager
ROOM_CHANNEL_PREFIX = "room:"
class RoomBroadcaster:
"""Cross-instance message fan-out (ARCHITECTURE.md phase 5).
Publishes to a per-room Redis channel; every app instance -- including
the one that published -- subscribes via a single pattern subscription
and forwards to its own locally connected WebSocket clients via
ConnectionManager. A single instance just talks to itself through Redis,
so there's no separate code path for the 1-instance vs N-instance case.
"""
def __init__(self, redis: Redis, manager: ConnectionManager) -> None:
self._redis = redis
self._manager = manager
async def publish(self, room_id: uuid.UUID, payload: dict) -> None:
await self._redis.publish(f"{ROOM_CHANNEL_PREFIX}{room_id}", json.dumps(payload))
async def listen(self) -> None:
pubsub = self._redis.pubsub()
await pubsub.psubscribe(f"{ROOM_CHANNEL_PREFIX}*")
try:
async for message in pubsub.listen():
if message["type"] != "pmessage":
continue
channel = message["channel"]
room_id = uuid.UUID(channel.removeprefix(ROOM_CHANNEL_PREFIX))
payload = json.loads(message["data"])
await self._manager.broadcast(room_id, payload)
finally:
await pubsub.punsubscribe(f"{ROOM_CHANNEL_PREFIX}*")
await pubsub.aclose()
+12 -6
View File
@@ -9,7 +9,7 @@ from app.database import get_db
from app.models import Room, RoomMembership, User from app.models import Room, RoomMembership, User
from app.services.message_service import create_message from app.services.message_service import create_message
from app.services.push_service import send_push_to_user from app.services.push_service import send_push_to_user
from app.ws.connection_manager import ConnectionManager from app.ws.presence import Presence
router = APIRouter(tags=["ws"]) router = APIRouter(tags=["ws"])
@@ -33,7 +33,7 @@ async def _is_room_member(db: AsyncSession, room_id: uuid.UUID, user_id: uuid.UU
async def _notify_offline_members( async def _notify_offline_members(
db: AsyncSession, db: AsyncSession,
manager: ConnectionManager, presence: Presence,
room_id: uuid.UUID, room_id: uuid.UUID,
sender: User, sender: User,
content: str, content: str,
@@ -42,7 +42,7 @@ async def _notify_offline_members(
select(RoomMembership.user_id).where(RoomMembership.room_id == room_id) select(RoomMembership.user_id).where(RoomMembership.room_id == room_id)
) )
member_ids = {row[0] for row in result.all()} member_ids = {row[0] for row in result.all()}
offline_ids = member_ids - manager.connected_user_ids(room_id) offline_ids = member_ids - await presence.connected_user_ids(room_id)
if not offline_ids: if not offline_ids:
return return
@@ -70,6 +70,8 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
await websocket.accept() await websocket.accept()
manager = websocket.app.state.connection_manager manager = websocket.app.state.connection_manager
presence: Presence = websocket.app.state.presence
broadcaster = websocket.app.state.broadcaster
joined_rooms: set[uuid.UUID] = set() joined_rooms: set[uuid.UUID] = set()
try: try:
@@ -90,7 +92,8 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
{"type": "error", "detail": "Not a member of this room"} {"type": "error", "detail": "Not a member of this room"}
) )
continue continue
manager.join(envelope.room_id, websocket, user.id) manager.join(envelope.room_id, websocket)
await presence.join(envelope.room_id, user.id)
joined_rooms.add(envelope.room_id) joined_rooms.add(envelope.room_id)
await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)}) await websocket.send_json({"type": "joined", "room_id": str(envelope.room_id)})
@@ -99,6 +102,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
await websocket.send_json({"type": "error", "detail": "room_id required"}) await websocket.send_json({"type": "error", "detail": "room_id required"})
continue continue
manager.leave(envelope.room_id, websocket) manager.leave(envelope.room_id, websocket)
await presence.leave(envelope.room_id, user.id)
joined_rooms.discard(envelope.room_id) joined_rooms.discard(envelope.room_id)
elif envelope.type == "message": elif envelope.type == "message":
@@ -115,7 +119,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
) )
continue continue
message = await create_message(db, envelope.room_id, user.id, envelope.content) message = await create_message(db, envelope.room_id, user.id, envelope.content)
await manager.broadcast( await broadcaster.publish(
envelope.room_id, envelope.room_id,
{ {
"type": "message", "type": "message",
@@ -128,7 +132,7 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
}, },
) )
await _notify_offline_members( await _notify_offline_members(
db, manager, envelope.room_id, user, envelope.content db, presence, envelope.room_id, user, envelope.content
) )
else: else:
@@ -140,3 +144,5 @@ async def chat_endpoint(websocket: WebSocket, db: AsyncSession = Depends(get_db)
pass pass
finally: finally:
manager.leave_all(websocket) manager.leave_all(websocket)
for room_id in joined_rooms:
await presence.leave(room_id, user.id)
+6 -18
View File
@@ -5,22 +5,19 @@ from fastapi import WebSocket
class ConnectionManager: class ConnectionManager:
"""In-memory, single-process WebSocket registry. """Local, single-process WebSocket socket registry.
Correct for a single app-server instance only; cross-instance fan-out via Purely about delivering to sockets connected to *this* process --
Redis pub/sub is a later phase (ARCHITECTURE.md phase 5). cross-instance fan-out lives in RoomBroadcaster, and cross-instance
"who's connected" for push lives in Presence, both backed by Redis
(ARCHITECTURE.md phase 5).
""" """
def __init__(self) -> None: def __init__(self) -> None:
self._rooms: dict[uuid.UUID, set[WebSocket]] = defaultdict(set) self._rooms: dict[uuid.UUID, set[WebSocket]] = defaultdict(set)
# A single connection can be joined to multiple rooms at once (one
# `join` message per room over the same socket), so this is keyed on
# the socket alone, not per-room.
self._ws_user: dict[WebSocket, uuid.UUID] = {}
def join(self, room_id: uuid.UUID, websocket: WebSocket, user_id: uuid.UUID) -> None: def join(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
self._rooms[room_id].add(websocket) self._rooms[room_id].add(websocket)
self._ws_user[websocket] = user_id
def leave(self, room_id: uuid.UUID, websocket: WebSocket) -> None: def leave(self, room_id: uuid.UUID, websocket: WebSocket) -> None:
self._rooms[room_id].discard(websocket) self._rooms[room_id].discard(websocket)
@@ -30,15 +27,6 @@ class ConnectionManager:
def leave_all(self, websocket: WebSocket) -> None: def leave_all(self, websocket: WebSocket) -> None:
for room_id in list(self._rooms.keys()): for room_id in list(self._rooms.keys()):
self.leave(room_id, websocket) self.leave(room_id, websocket)
self._ws_user.pop(websocket, None)
def connected_user_ids(self, room_id: uuid.UUID) -> set[uuid.UUID]:
"""Users (not just sockets) with an active connection to this room --
used to skip push notifications for anyone already watching, per
ARCHITECTURE.md's "members with no active connection" push flow."""
return {
self._ws_user[ws] for ws in self._rooms.get(room_id, ()) if ws in self._ws_user
}
async def broadcast(self, room_id: uuid.UUID, payload: dict) -> None: async def broadcast(self, room_id: uuid.UUID, payload: dict) -> None:
for websocket in list(self._rooms.get(room_id, ())): for websocket in list(self._rooms.get(room_id, ())):
+39
View File
@@ -0,0 +1,39 @@
import uuid
from redis.asyncio import Redis
class Presence:
"""Cross-instance "who's connected to this room," backed by a Redis hash
per room (field = user_id, value = connection refcount).
Refcounted rather than a plain set so a user with two connections to the
same room -- two tabs, or one per app instance -- doesn't get marked
offline when only one of those connections closes.
Known limitation: a hard crash (not a clean disconnect) leaks that
connection's increment forever, since there's no heartbeat/TTL here to
reclaim it -- out of scope for this phase, same category of
simplification as the "no server-side session revocation" note in the
README.
"""
def __init__(self, redis: Redis) -> None:
self._redis = redis
def _key(self, room_id: uuid.UUID) -> str:
return f"presence:{room_id}"
async def join(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
await self._redis.hincrby(self._key(room_id), str(user_id), 1)
async def leave(self, room_id: uuid.UUID, user_id: uuid.UUID) -> None:
key = self._key(room_id)
field = str(user_id)
remaining = await self._redis.hincrby(key, field, -1)
if remaining <= 0:
await self._redis.hdel(key, field)
async def connected_user_ids(self, room_id: uuid.UUID) -> set[uuid.UUID]:
fields = await self._redis.hkeys(self._key(room_id))
return {uuid.UUID(f) for f in fields}
+1
View File
@@ -15,6 +15,7 @@ dependencies = [
"argon2-cffi>=23.1", "argon2-cffi>=23.1",
"itsdangerous>=2.2", "itsdangerous>=2.2",
"pywebpush>=2.0", "pywebpush>=2.0",
"redis>=5.0",
] ]
[project.scripts] [project.scripts]
+38 -15
View File
@@ -1,3 +1,4 @@
import contextlib
import os import os
from pathlib import Path from pathlib import Path
@@ -6,6 +7,9 @@ os.environ.setdefault(
) )
os.environ.setdefault("SESSION_SECRET", "test-secret") os.environ.setdefault("SESSION_SECRET", "test-secret")
os.environ.setdefault("SESSION_HTTPS_ONLY", "false") os.environ.setdefault("SESSION_HTTPS_ONLY", "false")
# DB index 15 keeps test presence/pub-sub state separate from whatever a
# developer's local Redis is doing on db 0.
os.environ.setdefault("REDIS_URL", "redis://localhost:6379/15")
import pytest import pytest
import pytest_asyncio import pytest_asyncio
@@ -76,15 +80,15 @@ async def client(app):
@pytest.fixture @pytest.fixture
def ws_client(): def ws_client_factory():
# Starlette's TestClient (needed for websocket_connect, which httpx's # Starlette's TestClient (needed for websocket_connect, which httpx's
# async client doesn't support) runs the ASGI app on a background thread # async client doesn't support) runs the ASGI app on a background thread
# with its own event loop via anyio's BlockingPortal. asyncpg connections # with its own event loop via anyio's BlockingPortal. asyncpg connections
# are bound to the loop they're opened on, so this app gets its own # are bound to the loop they're opened on, so each app built here gets
# engine created here (no connections opened yet) rather than reusing # its own engine (no connections opened yet) rather than reusing the
# the `db_session`/`app` fixtures' engine, which belongs to pytest's # `db_session`/`app` fixtures' engine, which belongs to pytest's loop.
# loop. No per-test rollback here (see test_ws_chat.py for the # No per-test rollback here (see test_ws_chat.py for the unique-name
# unique-name convention that keeps tests independent without it). # convention that keeps tests independent without it).
# #
# poolclass=NullPool: with pooling, a WS test that does more than one # poolclass=NullPool: with pooling, a WS test that does more than one
# DB round trip per message (e.g. the offline-push lookup) can hit a # DB round trip per message (e.g. the offline-push lookup) can hit a
@@ -94,19 +98,38 @@ def ws_client():
# surfaces as "connection is closed". A fresh connection per session # surfaces as "connection is closed". A fresh connection per session
# sidesteps it; fine for tests, not something prod needs (prod isn't # sidesteps it; fine for tests, not something prod needs (prod isn't
# juggling a background portal thread against the main test thread). # juggling a background portal thread against the main test thread).
application = create_app() #
test_engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool) # A factory (not a single client) so tests can spin up more than one
test_session_factory = async_sessionmaker(test_engine, expire_on_commit=False) # independent app instance -- sharing the same test Postgres and Redis,
# like separate app-server processes behind Nginx would -- to exercise
# cross-instance broadcast/presence (see test_broadcast.py). Each
# TestClient is entered via an ExitStack so its lifespan (which opens
# the Redis connection/pubsub listener) starts immediately and all of
# them get torn down together at fixture teardown.
stack = contextlib.ExitStack()
async def _get_db(): def _make() -> TestClient:
async with test_session_factory() as session: application = create_app()
yield session test_engine = create_async_engine(TEST_DATABASE_URL, poolclass=NullPool)
test_session_factory = async_sessionmaker(test_engine, expire_on_commit=False)
application.dependency_overrides[get_db] = _get_db async def _get_db():
async with test_session_factory() as session:
yield session
with TestClient(application) as tc: application.dependency_overrides[get_db] = _get_db
tc = stack.enter_context(TestClient(application))
tc.session_factory = test_session_factory # type: ignore[attr-defined] tc.session_factory = test_session_factory # type: ignore[attr-defined]
yield tc return tc
yield _make
stack.close()
@pytest.fixture
def ws_client(ws_client_factory):
return ws_client_factory()
async def register_and_login( async def register_and_login(
+103
View File
@@ -0,0 +1,103 @@
import uuid
from app.schemas.user import UserCreate
from app.services.auth_service import register_user
def _unique(prefix: str) -> str:
return f"{prefix}-{uuid.uuid4().hex[:8]}"
def _register_ws(ws_client, username: str) -> dict:
async def _seed():
async with ws_client.session_factory() as session:
await register_user(
session,
UserCreate(username=username, email=f"{username}@example.com", password="password123"),
)
ws_client.portal.call(_seed)
resp = ws_client.post(
"/api/auth/login", json={"username_or_email": username, "password": "password123"}
)
assert resp.status_code == 200, resp.text
return resp.json()
def test_message_fans_out_across_instances(ws_client_factory):
# Two independent app instances -- separate ConnectionManager, separate
# Redis pubsub subscription, separate everything except the Postgres and
# Redis they're both pointed at -- the same way two app-server processes
# behind Nginx would be. Proves delivery actually crosses Redis, not
# just in-process delivery within a single ConnectionManager.
instance1 = ws_client_factory()
instance2 = ws_client_factory()
alice = _register_ws(instance1, _unique("alice"))
room = instance1.post("/api/rooms", json={"name": _unique("general")}).json()
bob = _register_ws(instance2, _unique("bob"))
instance2.post(f"/api/rooms/{room['id']}/join")
with instance2.websocket_connect("/ws/chat") as bob_ws:
bob_ws.send_json({"type": "join", "room_id": room["id"]})
assert bob_ws.receive_json()["type"] == "joined"
with instance1.websocket_connect("/ws/chat") as alice_ws:
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
alice_ws.send_json(
{"type": "message", "room_id": room["id"], "content": "hi from instance 1"}
)
assert alice_ws.receive_json()["type"] == "message"
received = bob_ws.receive_json()
assert received["type"] == "message"
assert received["content"] == "hi from instance 1"
assert received["username"] == alice["username"]
def test_presence_is_shared_across_instances(ws_client_factory, monkeypatch):
calls = []
monkeypatch.setattr("app.services.push_service.webpush", lambda **kw: calls.append(kw))
instance1 = ws_client_factory()
instance2 = ws_client_factory()
alice = _register_ws(instance1, _unique("alice"))
room = instance1.post("/api/rooms", json={"name": _unique("general")}).json()
bob = _register_ws(instance2, _unique("bob"))
instance2.post(f"/api/rooms/{room['id']}/join")
instance2.post(
"/api/push/subscribe",
json={
"endpoint": f"https://push.example.com/ep-{bob['id']}",
"keys": {"p256dh": "p256dh-bob", "auth": "auth-bob"},
},
)
with instance2.websocket_connect("/ws/chat") as bob_ws:
bob_ws.send_json({"type": "join", "room_id": room["id"]})
assert bob_ws.receive_json()["type"] == "joined"
with instance1.websocket_connect("/ws/chat") as alice_ws:
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
alice_ws.send_json({"type": "message", "room_id": room["id"], "content": "hi"})
assert alice_ws.receive_json()["type"] == "message"
# bob is connected -- just on the other instance -- so he should
# get the broadcast via Redis, not a push notification. If
# presence were still process-local (pre-phase-5 behavior) he'd
# look offline to instance1 and get a redundant push.
assert bob_ws.receive_json()["type"] == "message"
# Sync barrier: the handler processes frames strictly
# sequentially, so a second (idempotent) join only acks once the
# "message" frame's full handling -- including the offline-push
# step -- has completed on instance1.
alice_ws.send_json({"type": "join", "room_id": room["id"]})
assert alice_ws.receive_json()["type"] == "joined"
assert calls == []