feat(B-WS): WebSocket Helpers + Redis Pub/Sub + Error-Handling
Check Cross-Plugin Imports / check (push) Has been cancelled
Check Cross-Plugin Imports / check (push) Has been cancelled
B-WS: app/core/ws_helpers.py (NEU) — gemeinsame WebSocket Helpers - authenticate_ws: Session-Auth für WebSocket (Cookie/Token → User/Tenant) - check_ws_origin: Origin-Check (delegiert auf verify_ws_origin) - check_ws_tenant: User-Tenant-Membership-Check - cleanup_ws_connection: Connection aus Registry entfernen + WS schließen - start_heartbeat: Background Ping-Task - send_ws_error: strukturierte Error-Message an Client - handle_ws_message: Message-Dispatch mit Error-Handling B-WS: app/core/ws_pubsub.py (NEU) — Redis Pub/Sub für Multi-Worker-Fanout - publish_to_channel / subscribe_to_channel - get_tenant_channel / broadcast_to_tenants B-WS: WebSocketManager + AIUIControlWSManager angepasst - connect() nutzt authenticate_ws + check_ws_origin + check_ws_tenant - disconnect() nutzt cleanup_ws_connection + cancelt Heartbeat/PubSub - broadcast() unterstützt Redis Pub/Sub Fanout B-ERR-WS: WS Error-Handling in ws_helpers integriert - send_ws_error für strukturierte Errors - handle_ws_message fängt Handler-Exceptions B-WS-TEST: 24 Tests in test_ws_helpers.py — alle grün - Auth, Origin, Error, Dispatch, Cleanup, Heartbeat, Pub/Sub Roundtrip - Keine Regression: 47/47 Resilience+Hooks Tests grün
This commit is contained in:
@@ -0,0 +1,500 @@
|
||||
"""Tests for WebSocket helpers (ws_helpers) and Redis Pub/Sub (ws_pubsub).
|
||||
|
||||
Covers B-WS-TEST + B-ERR-WS-TEST requirements.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from app.core.ws_helpers import (
|
||||
authenticate_ws,
|
||||
check_ws_origin,
|
||||
cleanup_ws_connection,
|
||||
handle_ws_message,
|
||||
send_ws_error,
|
||||
start_heartbeat,
|
||||
)
|
||||
from app.core.ws_pubsub import (
|
||||
broadcast_to_tenants,
|
||||
get_tenant_channel,
|
||||
publish_to_channel,
|
||||
subscribe_to_channel,
|
||||
)
|
||||
|
||||
|
||||
# ─── Mock WebSocket ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class MockWebSocket:
|
||||
"""Minimal mock WebSocket for testing WS helpers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cookies: dict[str, str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
query_params: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self.cookies = cookies or {}
|
||||
self.headers = headers or {}
|
||||
self.query_params = query_params or {}
|
||||
self._closed = False
|
||||
self._close_code: int | None = None
|
||||
self._close_reason: str | None = None
|
||||
self._sent: list[str] = []
|
||||
self._accepted = False
|
||||
|
||||
async def accept(self) -> None:
|
||||
self._accepted = True
|
||||
|
||||
async def close(self, code: int = 1000, reason: str = "") -> None:
|
||||
self._closed = True
|
||||
self._close_code = code
|
||||
self._close_reason = reason
|
||||
|
||||
async def send_text(self, text: str) -> None:
|
||||
if self._closed:
|
||||
raise RuntimeError("WebSocket is closed")
|
||||
self._sent.append(text)
|
||||
|
||||
@property
|
||||
def sent_messages(self) -> list[dict]:
|
||||
return [json.loads(t) for t in self._sent]
|
||||
|
||||
|
||||
# ─── authenticate_ws ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestAuthenticateWs:
|
||||
"""Tests for authenticate_ws."""
|
||||
|
||||
async def test_authenticate_ws_valid_session(self, db_session, redis_client):
|
||||
"""authenticate_ws returns user info when session is valid."""
|
||||
from app.core.auth import create_session, hash_password
|
||||
from app.models.tenant import Tenant
|
||||
from app.models.user import User, UserTenant
|
||||
from app.models.role import Role
|
||||
|
||||
# Create tenant, user, role, membership
|
||||
tenant = Tenant(name="Test Tenant", slug="test-tenant")
|
||||
db_session.add(tenant)
|
||||
await db_session.flush()
|
||||
|
||||
user = User(
|
||||
email="wsauth@test.com",
|
||||
name="WS Auth Test",
|
||||
password_hash=hash_password("TestPass123!"),
|
||||
is_active=True,
|
||||
preferences={},
|
||||
)
|
||||
db_session.add(user)
|
||||
await db_session.flush()
|
||||
|
||||
role = Role(
|
||||
tenant_id=tenant.id,
|
||||
name="admin",
|
||||
permissions={"*": {"*": True}},
|
||||
denied_permissions=[],
|
||||
field_permissions={},
|
||||
)
|
||||
db_session.add(role)
|
||||
await db_session.flush()
|
||||
|
||||
ut = UserTenant(
|
||||
user_id=user.id,
|
||||
tenant_id=tenant.id,
|
||||
is_default=True,
|
||||
role="admin",
|
||||
role_id=role.id,
|
||||
)
|
||||
db_session.add(ut)
|
||||
await db_session.flush()
|
||||
await db_session.commit()
|
||||
|
||||
# Create session
|
||||
session_id, csrf_token = await create_session(
|
||||
db_session, redis_client, user, tenant.id, role="admin"
|
||||
)
|
||||
|
||||
# Create mock WebSocket with session cookie
|
||||
ws = MockWebSocket(cookies={"leocrm_session": session_id})
|
||||
|
||||
auth = await authenticate_ws(ws, db_session)
|
||||
assert auth is not None
|
||||
assert auth["user_id"] == str(user.id)
|
||||
assert auth["tenant_id"] == str(tenant.id)
|
||||
assert auth["session_id"] == session_id
|
||||
assert auth["role"] == "admin"
|
||||
assert ws._closed is False
|
||||
|
||||
async def test_authenticate_ws_missing_cookie(self, db_session):
|
||||
"""authenticate_ws closes WS and returns None when no session cookie."""
|
||||
ws = MockWebSocket(cookies={})
|
||||
|
||||
auth = await authenticate_ws(ws, db_session)
|
||||
assert auth is None
|
||||
assert ws._closed is True
|
||||
assert ws._close_code == 4401
|
||||
|
||||
async def test_authenticate_ws_invalid_session(self, db_session):
|
||||
"""authenticate_ws closes WS and returns None when session is invalid."""
|
||||
ws = MockWebSocket(cookies={"leocrm_session": "nonexistent-session-id"})
|
||||
|
||||
auth = await authenticate_ws(ws, db_session)
|
||||
assert auth is None
|
||||
assert ws._closed is True
|
||||
assert ws._close_code == 4401
|
||||
|
||||
|
||||
# ─── check_ws_origin ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestCheckWsOrigin:
|
||||
"""Tests for check_ws_origin."""
|
||||
|
||||
async def test_check_ws_origin_valid(self):
|
||||
"""check_ws_origin returns True for valid origin + CSRF."""
|
||||
from app.core.auth import create_session, hash_password
|
||||
from app.models.tenant import Tenant
|
||||
from app.models.user import User, UserTenant
|
||||
from app.models.role import Role
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
|
||||
# We need a full setup with session in Redis for CSRF validation
|
||||
# Use patch to mock verify_ws_origin returning True
|
||||
with patch("app.core.ws_helpers.verify_ws_origin", return_value=True):
|
||||
ws = MockWebSocket()
|
||||
result = await check_ws_origin(ws)
|
||||
assert result is True
|
||||
assert ws._closed is False
|
||||
|
||||
async def test_check_ws_origin_invalid(self):
|
||||
"""check_ws_origin closes WS and returns False for invalid origin."""
|
||||
with patch("app.core.ws_helpers.verify_ws_origin", return_value=False):
|
||||
ws = MockWebSocket()
|
||||
result = await check_ws_origin(ws)
|
||||
assert result is False
|
||||
assert ws._closed is True
|
||||
assert ws._close_code == 4403
|
||||
|
||||
|
||||
# ─── send_ws_error ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestSendWsError:
|
||||
"""Tests for send_ws_error."""
|
||||
|
||||
async def test_send_ws_error_basic(self):
|
||||
"""send_ws_error sends structured error message."""
|
||||
ws = MockWebSocket()
|
||||
await send_ws_error(ws, "test_error", "Something went wrong")
|
||||
|
||||
assert len(ws.sent_messages) == 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["type"] == "error"
|
||||
assert msg["code"] == "test_error"
|
||||
assert msg["detail"] == "Something went wrong"
|
||||
assert "trace_id" not in msg
|
||||
|
||||
async def test_send_ws_error_with_trace_id(self):
|
||||
"""send_ws_error includes trace_id when provided."""
|
||||
ws = MockWebSocket()
|
||||
await send_ws_error(ws, "test_error", "Failed", trace_id="trace-123")
|
||||
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["trace_id"] == "trace-123"
|
||||
|
||||
async def test_send_ws_error_on_closed_ws(self):
|
||||
"""send_ws_error does not raise when WS is closed."""
|
||||
ws = MockWebSocket()
|
||||
ws._closed = True
|
||||
# Should not raise
|
||||
await send_ws_error(ws, "closed_error", "WS already closed")
|
||||
|
||||
|
||||
# ─── handle_ws_message ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestHandleWsMessage:
|
||||
"""Tests for handle_ws_message."""
|
||||
|
||||
async def test_handle_ws_message_dispatch(self):
|
||||
"""handle_ws_message dispatches to correct handler."""
|
||||
ws = MockWebSocket()
|
||||
called = False
|
||||
|
||||
async def ping_handler(websocket, msg):
|
||||
nonlocal called
|
||||
called = True
|
||||
assert msg["type"] == "ping"
|
||||
|
||||
handlers = {"ping": ping_handler}
|
||||
message = json.dumps({"type": "ping"})
|
||||
await handle_ws_message(ws, message, handlers)
|
||||
|
||||
assert called is True
|
||||
|
||||
async def test_handle_ws_message_unknown_type(self):
|
||||
"""handle_ws_message sends error for unknown type."""
|
||||
ws = MockWebSocket()
|
||||
handlers = {"ping": AsyncMock()}
|
||||
message = json.dumps({"type": "unknown_type"})
|
||||
await handle_ws_message(ws, message, handlers)
|
||||
|
||||
assert len(ws.sent_messages) == 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["type"] == "error"
|
||||
assert msg["code"] == "unknown_type"
|
||||
|
||||
async def test_handle_ws_message_handler_exception(self):
|
||||
"""handle_ws_message sends error when handler raises exception."""
|
||||
ws = MockWebSocket()
|
||||
|
||||
async def bad_handler(websocket, msg):
|
||||
raise ValueError("Handler crashed")
|
||||
|
||||
handlers = {"ping": bad_handler}
|
||||
message = json.dumps({"type": "ping"})
|
||||
await handle_ws_message(ws, message, handlers)
|
||||
|
||||
assert len(ws.sent_messages) == 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["type"] == "error"
|
||||
assert msg["code"] == "handler_error"
|
||||
assert "Handler crashed" in msg["detail"]
|
||||
|
||||
async def test_handle_ws_message_invalid_json(self):
|
||||
"""handle_ws_message sends error for invalid JSON."""
|
||||
ws = MockWebSocket()
|
||||
handlers = {}
|
||||
await handle_ws_message(ws, "not json at all", handlers)
|
||||
|
||||
assert len(ws.sent_messages) == 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["code"] == "invalid_json"
|
||||
|
||||
async def test_handle_ws_message_missing_type(self):
|
||||
"""handle_ws_message sends error when type field is missing."""
|
||||
ws = MockWebSocket()
|
||||
handlers = {}
|
||||
message = json.dumps({"data": "no type here"})
|
||||
await handle_ws_message(ws, message, handlers)
|
||||
|
||||
assert len(ws.sent_messages) == 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["code"] == "missing_type"
|
||||
|
||||
|
||||
# ─── cleanup_ws_connection ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestCleanupWsConnection:
|
||||
"""Tests for cleanup_ws_connection."""
|
||||
|
||||
async def test_cleanup_removes_from_registry(self):
|
||||
"""cleanup_ws_connection removes WebSocket from registry."""
|
||||
ws1 = MockWebSocket()
|
||||
ws2 = MockWebSocket()
|
||||
registry: dict[str, list] = {"user1": [ws1, ws2]}
|
||||
|
||||
await cleanup_ws_connection(ws1, "user1", registry)
|
||||
|
||||
assert "user1" in registry
|
||||
assert ws1 not in registry["user1"]
|
||||
assert ws2 in registry["user1"]
|
||||
assert len(registry["user1"]) == 1
|
||||
|
||||
async def test_cleanup_removes_empty_user(self):
|
||||
"""cleanup_ws_connection removes user entry when no connections left."""
|
||||
ws = MockWebSocket()
|
||||
registry: dict[str, list] = {"user1": [ws]}
|
||||
|
||||
await cleanup_ws_connection(ws, "user1", registry)
|
||||
|
||||
assert "user1" not in registry
|
||||
|
||||
async def test_cleanup_unknown_user(self):
|
||||
"""cleanup_ws_connection handles unknown user gracefully."""
|
||||
ws = MockWebSocket()
|
||||
registry: dict[str, list] = {}
|
||||
|
||||
await cleanup_ws_connection(ws, "unknown_user", registry)
|
||||
assert "unknown_user" not in registry
|
||||
|
||||
|
||||
# ─── start_heartbeat ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestStartHeartbeat:
|
||||
"""Tests for start_heartbeat."""
|
||||
|
||||
async def test_heartbeat_sends_ping(self):
|
||||
"""start_heartbeat sends ping messages at interval."""
|
||||
ws = MockWebSocket()
|
||||
task = await start_heartbeat(ws, interval=0)
|
||||
|
||||
# Wait a tiny bit for the first ping
|
||||
await asyncio.sleep(0.05)
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert len(ws._sent) >= 1
|
||||
msg = ws.sent_messages[0]
|
||||
assert msg["type"] == "ping"
|
||||
|
||||
async def test_heartbeat_cancellable(self):
|
||||
"""start_heartbeat task can be cancelled."""
|
||||
ws = MockWebSocket()
|
||||
task = await start_heartbeat(ws, interval=10)
|
||||
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
assert task.cancelled() or task.done()
|
||||
|
||||
|
||||
# ─── get_tenant_channel ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGetTenantChannel:
|
||||
"""Tests for get_tenant_channel."""
|
||||
|
||||
def test_channel_naming(self):
|
||||
"""get_tenant_channel returns ws:{tenant_id}:{topic}."""
|
||||
tenant_id = uuid.uuid4()
|
||||
channel = get_tenant_channel(tenant_id, "kommunikation")
|
||||
assert channel == f"ws:{tenant_id}:kommunikation"
|
||||
|
||||
def test_channel_naming_different_topic(self):
|
||||
"""get_tenant_channel works with different topics."""
|
||||
tenant_id = uuid.uuid4()
|
||||
channel = get_tenant_channel(tenant_id, "ai_ui_control")
|
||||
assert channel == f"ws:{tenant_id}:ai_ui_control"
|
||||
|
||||
|
||||
# ─── publish_to_channel + subscribe_to_channel ──────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestPubSub:
|
||||
"""Tests for Redis Pub/Sub helpers."""
|
||||
|
||||
async def test_publish_subscribe_roundtrip(self, redis_client):
|
||||
"""publish_to_channel + subscribe_to_channel roundtrip."""
|
||||
received: list[dict] = []
|
||||
channel = f"test-roundtrip-{uuid.uuid4()}"
|
||||
event = asyncio.Event()
|
||||
|
||||
async def handler(msg: dict) -> None:
|
||||
received.append(msg)
|
||||
event.set()
|
||||
|
||||
# Subscribe
|
||||
task = await subscribe_to_channel(channel, handler)
|
||||
|
||||
# Give subscriber a moment to connect
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Publish
|
||||
test_msg = {"type": "test", "data": "hello"}
|
||||
await publish_to_channel(channel, test_msg)
|
||||
|
||||
# Wait for message
|
||||
await asyncio.wait_for(event.wait(), timeout=2.0)
|
||||
|
||||
assert len(received) == 1
|
||||
assert received[0]["type"] == "test"
|
||||
assert received[0]["data"] == "hello"
|
||||
|
||||
# Cleanup
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def test_broadcast_to_tenants(self, redis_client):
|
||||
"""broadcast_to_tenants publishes to the correct tenant channel."""
|
||||
tenant_id = uuid.uuid4()
|
||||
received: list[dict] = []
|
||||
event = asyncio.Event()
|
||||
|
||||
channel = get_tenant_channel(tenant_id, "test_topic")
|
||||
|
||||
async def handler(msg: dict) -> None:
|
||||
received.append(msg)
|
||||
event.set()
|
||||
|
||||
task = await subscribe_to_channel(channel, handler)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
test_msg = {"type": "broadcast", "content": "tenant message"}
|
||||
await broadcast_to_tenants(tenant_id, "test_topic", test_msg)
|
||||
|
||||
await asyncio.wait_for(event.wait(), timeout=2.0)
|
||||
|
||||
assert len(received) == 1
|
||||
assert received[0]["type"] == "broadcast"
|
||||
assert received[0]["content"] == "tenant message"
|
||||
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def test_publish_to_channel_no_subscribers(self, redis_client):
|
||||
"""publish_to_channel works even with no subscribers."""
|
||||
channel = f"test-no-sub-{uuid.uuid4()}"
|
||||
await publish_to_channel(channel, {"type": "noop"})
|
||||
# Should not raise
|
||||
|
||||
async def test_subscribe_multiple_messages(self, redis_client):
|
||||
"""subscribe_to_channel handler receives multiple messages."""
|
||||
received: list[dict] = []
|
||||
channel = f"test-multi-{uuid.uuid4()}"
|
||||
count_event = asyncio.Event()
|
||||
msg_count = 0
|
||||
|
||||
async def handler(msg: dict) -> None:
|
||||
nonlocal msg_count
|
||||
received.append(msg)
|
||||
msg_count += 1
|
||||
if msg_count >= 3:
|
||||
count_event.set()
|
||||
|
||||
task = await subscribe_to_channel(channel, handler)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
for i in range(3):
|
||||
await publish_to_channel(channel, {"type": "msg", "index": i})
|
||||
|
||||
await asyncio.wait_for(count_event.wait(), timeout=3.0)
|
||||
|
||||
assert len(received) == 3
|
||||
assert received[0]["index"] == 0
|
||||
assert received[2]["index"] == 2
|
||||
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
Reference in New Issue
Block a user