mjalmoudy94
f1e5df29ce
- Backend: ConnectionManager, Redis-backed bus with in-process fallback, channel authorization, and typed emit helpers wired into workers (ai_runtime, sandbox_gateway), workforce/sandbox/billing/audit/error services. - Frontend: single shared WS client with backoff + heartbeat, RealtimeProvider, useLiveChannel/useLiveEvent hooks, live indicators on task detail, manager chat, dashboard, traces, sandbox terminal, admin errors, and billing. - Contracts: WsEvent/WsServerFrame/WsClientFrame types mirroring the port. All 255 backend tests pass; web typecheck and build pass.
417 lines
15 KiB
Python
417 lines
15 KiB
Python
"""Realtime push layer: authorization, fan-out, and the live socket end-to-end.
|
|
|
|
Three levels, mirroring the package's structure:
|
|
|
|
* **Authorization** — a channel is only ever granted through
|
|
``authorize_channel``; every rule is tested for both the allow and the deny
|
|
side, because a rule that always denies is indistinguishable from a bug.
|
|
* **Fan-out** — the manager delivers to the subscribers of a channel, keeps
|
|
ordering, de-duplicates relayed events and cleans up on unsubscribe.
|
|
* **The socket** — a real WebSocket against the real app: subscribe, act over
|
|
HTTP, and assert the event arrives without a reload.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import suppress
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
import app.infrastructure.db.models as _orm_models # noqa: F401 - registers metadata
|
|
from app.application.ports.realtime import RealtimeEvent, set_publisher
|
|
from app.core.config import settings
|
|
from app.core.database import Base, get_session
|
|
from app.core.security import create_access_token
|
|
from app.domain.base import new_id
|
|
from app.infrastructure.db.models import UserRow
|
|
from app.interfaces.realtime.authz import authorize_channel
|
|
from app.interfaces.realtime.bus import get_bus
|
|
from app.interfaces.realtime.manager import Connection, ConnectionManager
|
|
from app.main import app
|
|
|
|
API_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class FakeSocket:
|
|
"""The manager only needs ``send_text`` and a connected ``client_state``.
|
|
|
|
Both matter: the sender task checks ``client_state`` before writing, so a
|
|
fake without it would silently swallow every frame and make the manager
|
|
look broken.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.sent: list[str] = []
|
|
from starlette.websockets import WebSocketState
|
|
|
|
self.client_state = WebSocketState.CONNECTED
|
|
|
|
async def send_text(self, payload: str) -> None:
|
|
self.sent.append(payload)
|
|
|
|
|
|
def _event(event_id: str, *channels: str) -> RealtimeEvent:
|
|
return RealtimeEvent(type="test.event", channels=channels, data={"id": event_id}, id=event_id)
|
|
|
|
|
|
async def _allow(_: Connection, __: str) -> bool:
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Authorization
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _user_with_org(session: AsyncSession, *, platform_owner: bool = False):
|
|
"""A user, and optionally an organization they belong to."""
|
|
user_id = new_id()
|
|
session.add(
|
|
UserRow(
|
|
id=user_id,
|
|
email=f"{user_id}@example.com",
|
|
password_hash="x",
|
|
is_platform_owner=platform_owner,
|
|
)
|
|
)
|
|
await session.flush()
|
|
return user_id
|
|
|
|
|
|
async def _org_with_member(session: AsyncSession, user_id: str) -> str:
|
|
from app.domain.identity.models import Membership, Organization, Role
|
|
from app.infrastructure.repositories.identity import (
|
|
MembershipRepository,
|
|
OrganizationRepository,
|
|
)
|
|
|
|
organization = await OrganizationRepository(session).create(
|
|
Organization(name="سازمان تست", slug=new_id()[:8], owner_user_id=user_id)
|
|
)
|
|
await MembershipRepository(session).add(
|
|
Membership(organization_id=organization.id, user_id=user_id, role=Role.ADMIN)
|
|
)
|
|
return organization.id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_member_may_subscribe_to_own_org(session_factory):
|
|
async with session_factory() as session:
|
|
user_id = await _user_with_org(session)
|
|
org_id = await _org_with_member(session, user_id)
|
|
allowed = await authorize_channel(
|
|
session, user_id=user_id, is_platform_owner=False, channel=f"org:{org_id}"
|
|
)
|
|
assert allowed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_outsider_may_not_subscribe_to_foreign_org(session_factory):
|
|
async with session_factory() as session:
|
|
insider = await _user_with_org(session)
|
|
outsider = await _user_with_org(session)
|
|
org_id = await _org_with_member(session, insider)
|
|
allowed = await authorize_channel(
|
|
session, user_id=outsider, is_platform_owner=False, channel=f"org:{org_id}"
|
|
)
|
|
assert allowed is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unknown_channel_kinds_are_denied(session_factory):
|
|
async with session_factory() as session:
|
|
user_id = await _user_with_org(session)
|
|
allowed = await authorize_channel(
|
|
session, user_id=user_id, is_platform_owner=False, channel="nonsense:whatever"
|
|
)
|
|
assert allowed is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_admin_channel_needs_platform_owner(session_factory):
|
|
async with session_factory() as session:
|
|
owner = await _user_with_org(session, platform_owner=True)
|
|
member = await _user_with_org(session)
|
|
owner_ok = await authorize_channel(
|
|
session, user_id=owner, is_platform_owner=True, channel="admin"
|
|
)
|
|
member_ok = await authorize_channel(
|
|
session, user_id=member, is_platform_owner=False, channel="admin"
|
|
)
|
|
assert owner_ok is True
|
|
assert member_ok is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_channel_is_private(session_factory):
|
|
async with session_factory() as session:
|
|
user_id = await _user_with_org(session)
|
|
other = await _user_with_org(session)
|
|
own = await authorize_channel(
|
|
session, user_id=user_id, is_platform_owner=False, channel=f"user:{user_id}"
|
|
)
|
|
foreign = await authorize_channel(
|
|
session, user_id=other, is_platform_owner=False, channel=f"user:{user_id}"
|
|
)
|
|
assert own is True
|
|
assert foreign is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fan-out
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_delivers_to_subscribers_only():
|
|
manager = ConnectionManager()
|
|
listening = FakeSocket()
|
|
other = FakeSocket()
|
|
|
|
subscribed = Connection(user_id="u1", is_platform_owner=False, websocket=listening)
|
|
stranger = Connection(user_id="u2", is_platform_owner=False, websocket=other)
|
|
for connection in (subscribed, stranger):
|
|
await manager.register(connection)
|
|
await manager.subscribe(subscribed, ["org:1"], _allow)
|
|
|
|
delivered = await manager.dispatch(_event("e1", "org:1"))
|
|
assert delivered == 1
|
|
|
|
# The write happens on the connection's own sender task; give the loop the
|
|
# turn it needs to flush the queue.
|
|
await asyncio.sleep(0.05)
|
|
assert len(listening.sent) == 1
|
|
assert other.sent == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_wraps_events_in_the_wire_envelope():
|
|
"""A browser sees one discriminator: ``{"type": "event", "event": {...}}``."""
|
|
import json
|
|
|
|
manager = ConnectionManager()
|
|
socket = FakeSocket()
|
|
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
|
await manager.register(connection)
|
|
await manager.subscribe(connection, ["org:1"], _allow)
|
|
await manager.dispatch(_event("e1", "org:1"))
|
|
|
|
await asyncio.sleep(0.05)
|
|
assert len(socket.sent) == 1
|
|
frame = json.loads(socket.sent[0])
|
|
assert frame["type"] == "event"
|
|
assert frame["event"]["id"] == "e1"
|
|
assert frame["event"]["type"] == "test.event"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_drops_duplicate_events():
|
|
"""In-process publish plus a relayed copy deliver exactly once."""
|
|
manager = ConnectionManager()
|
|
socket = FakeSocket()
|
|
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
|
await manager.register(connection)
|
|
await manager.subscribe(connection, ["org:1"], _allow)
|
|
|
|
event = _event("dup", "org:1")
|
|
first = await manager.dispatch(event)
|
|
second = await manager.dispatch(event)
|
|
assert (first, second) == (1, 0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unsubscribe_stops_delivery():
|
|
manager = ConnectionManager()
|
|
socket = FakeSocket()
|
|
connection = Connection(user_id="u1", is_platform_owner=False, websocket=socket)
|
|
await manager.register(connection)
|
|
await manager.subscribe(connection, ["org:1", "org:2"], _allow)
|
|
await manager.unsubscribe(connection, ["org:1"])
|
|
|
|
assert await manager.dispatch(_event("e1", "org:1")) == 0
|
|
assert await manager.dispatch(_event("e2", "org:2")) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_publisher_never_raises_on_bad_listener():
|
|
"""A listener that explodes must not break the publisher or its peers."""
|
|
|
|
async def exploding(event: RealtimeEvent) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
bus = get_bus()
|
|
bus.attach(exploding)
|
|
try:
|
|
set_publisher(bus)
|
|
# Must not raise despite the exploding listener.
|
|
from app.application.ports.realtime import publish
|
|
|
|
await publish(_event("e1", "org:1"))
|
|
finally:
|
|
bus.detach(exploding)
|
|
set_publisher(_NullPublisher())
|
|
|
|
|
|
class _NullPublisher:
|
|
async def publish(self, event: RealtimeEvent) -> None:
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The socket, end to end
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _create_schema(engine) -> None:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
@pytest.fixture()
|
|
def ws_client(monkeypatch):
|
|
"""A TestClient whose API and socket both talk to one throwaway database.
|
|
|
|
File-based SQLite with ``NullPool`` on purpose: the app runs on the test
|
|
portal's event loop while the schema is created from the test's own, and a
|
|
pooled connection cannot cross that boundary. The database file lives in
|
|
the package directory — ``*.db`` is gitignored — so the fixture needs no
|
|
writable temp directory outside the workspace.
|
|
"""
|
|
db_path = API_ROOT / f"realtime-test-{new_id()}.db"
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", poolclass=NullPool)
|
|
asyncio.run(_create_schema(engine))
|
|
factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
monkeypatch.setattr("app.main.engine", engine)
|
|
monkeypatch.setattr("app.interfaces.realtime.router.SessionFactory", factory)
|
|
monkeypatch.setattr(settings, "realtime_transport", "local")
|
|
|
|
async def _override_session():
|
|
async with factory() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
|
|
app.dependency_overrides[get_session] = _override_session
|
|
# The manager is module state; start each run from a clean slate so a
|
|
# connection from another test cannot colour this one's fan-out.
|
|
from app.interfaces.realtime.router import realtime_manager
|
|
|
|
realtime_manager()._connections.clear()
|
|
realtime_manager()._by_channel.clear()
|
|
try:
|
|
with TestClient(app) as client:
|
|
yield client
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
asyncio.run(engine.dispose())
|
|
with suppress(OSError):
|
|
db_path.unlink()
|
|
|
|
|
|
def _frames_until(websocket, wanted_type: str, *, max_frames: int = 20) -> dict:
|
|
"""Read frames until ``wanted_type`` arrives — hello/subscribed come first."""
|
|
for _ in range(max_frames):
|
|
frame = websocket.receive_json()
|
|
if frame.get("type") == "event" and frame.get("event", {}).get("type") == wanted_type:
|
|
return frame["event"]
|
|
raise AssertionError(f"no {wanted_type} event arrived")
|
|
|
|
|
|
def test_socket_rejects_a_missing_token(ws_client):
|
|
with pytest.raises(Exception), ws_client.websocket_connect("/api/ws"): # noqa: B017, SIM117
|
|
# No token: the endpoint closes with 1008 before accepting.
|
|
pass
|
|
|
|
|
|
def test_task_creation_is_pushed_live(ws_client):
|
|
# Sign up, create an organization — the tenant the socket subscribes to.
|
|
signup = ws_client.post(
|
|
"/api/auth/signup",
|
|
json={"email": "live@example.com", "password": "password123", "name": "زنده"},
|
|
)
|
|
assert signup.status_code == 201, signup.text
|
|
token = signup.json()["access_token"]
|
|
org = ws_client.post(
|
|
"/api/organizations",
|
|
json={"name": "سازمان زنده"},
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert org.status_code == 201, org.text
|
|
org_id = org.json()["id"]
|
|
|
|
with ws_client.websocket_connect(f"/api/ws?token={token}") as socket:
|
|
hello = socket.receive_json()
|
|
assert hello["type"] == "hello"
|
|
|
|
socket.send_json({"action": "subscribe", "channels": [f"org:{org_id}"]})
|
|
subscribed = socket.receive_json()
|
|
assert subscribed["type"] == "subscribed"
|
|
assert subscribed["accepted"] == [f"org:{org_id}"]
|
|
|
|
# A foreign channel is refused, not silently ignored.
|
|
socket.send_json({"action": "subscribe", "channels": ["org:does-not-exist"]})
|
|
refused = socket.receive_json()
|
|
assert refused["type"] == "subscribed"
|
|
assert refused["rejected"] == ["org:does-not-exist"]
|
|
|
|
created = ws_client.post(
|
|
f"/api/organizations/{org_id}/tasks",
|
|
json={"title": "کار زنده"},
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert created.status_code == 201, created.text
|
|
|
|
event = _frames_until(socket, "task.created")
|
|
assert event["data"]["title"] == "کار زنده"
|
|
assert event["organization_id"] == org_id
|
|
|
|
|
|
def test_second_user_cannot_read_another_org(ws_client):
|
|
"""Two tenants, one socket each: an outsider's org is out of reach."""
|
|
tokens = []
|
|
org_ids = []
|
|
for index in range(2):
|
|
signup = ws_client.post(
|
|
"/api/auth/signup",
|
|
json={"email": f"user{index}@example.com", "password": "password123"},
|
|
)
|
|
assert signup.status_code == 201, signup.text
|
|
token = signup.json()["access_token"]
|
|
org = ws_client.post(
|
|
"/api/organizations",
|
|
json={"name": f"org {index}"},
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert org.status_code == 201, org.text
|
|
tokens.append(token)
|
|
org_ids.append(org.json()["id"])
|
|
|
|
with ws_client.websocket_connect(f"/api/ws?token={tokens[0]}") as socket:
|
|
socket.receive_json() # hello
|
|
socket.send_json({"action": "subscribe", "channels": [f"org:{org_ids[1]}"]})
|
|
refused = socket.receive_json()
|
|
assert refused["rejected"] == [f"org:{org_ids[1]}"]
|
|
|
|
|
|
def test_socket_token_identifies_a_real_user(ws_client):
|
|
"""A well-formed token for a deleted user must not open a socket."""
|
|
user_id = new_id()
|
|
token = create_access_token(user_id)
|
|
with pytest.raises(Exception), ws_client.websocket_connect(f"/api/ws?token={token}"): # noqa: B017, SIM117
|
|
# Closed with 1008: the token is well-formed but names no user.
|
|
pass
|