388 lines
14 KiB
Python
388 lines
14 KiB
Python
# tests/test_collab.py — Tests for real-time collaboration (ROADMAP #62)
|
|
"""Unit and integration tests for the collaboration WebSocket relay."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from backend.collab import (
|
|
CollabManager,
|
|
CollabRoom,
|
|
authenticate_websocket,
|
|
color_for_index,
|
|
)
|
|
|
|
|
|
class _StubWebSocket:
|
|
"""Minimal WebSocket stub for auth unit tests."""
|
|
|
|
def __init__(self, query=None, cookies=None):
|
|
self.query_params = query or {}
|
|
self.cookies = cookies or {}
|
|
|
|
|
|
def _recv_json(ws) -> dict:
|
|
return json.loads(ws.receive_text())
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_collab_manager():
|
|
"""Isolate the collaboration singleton between tests."""
|
|
from backend.collab import collab_manager
|
|
|
|
collab_manager._rooms.clear()
|
|
collab_manager._next_conn_id = 1
|
|
yield
|
|
collab_manager._rooms.clear()
|
|
collab_manager._next_conn_id = 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pure helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_color_for_index_is_deterministic_and_cycles():
|
|
assert color_for_index(0) == color_for_index(10)
|
|
assert color_for_index(0) != color_for_index(1)
|
|
assert color_for_index(0).startswith("#")
|
|
|
|
|
|
def test_authenticate_websocket_disabled_returns_anonymous(monkeypatch):
|
|
monkeypatch.setenv("OBSIGATE_AUTH_ENABLED", "false")
|
|
user = authenticate_websocket(_StubWebSocket())
|
|
assert user is not None
|
|
assert user["role"] == "admin"
|
|
assert user["vaults"] == ["*"]
|
|
|
|
|
|
def test_authenticate_websocket_no_token_returns_none(monkeypatch):
|
|
monkeypatch.setenv("OBSIGATE_AUTH_ENABLED", "true")
|
|
assert authenticate_websocket(_StubWebSocket()) is None
|
|
|
|
|
|
def test_authenticate_websocket_invalid_token_returns_none(monkeypatch):
|
|
monkeypatch.setenv("OBSIGATE_AUTH_ENABLED", "true")
|
|
assert authenticate_websocket(_StubWebSocket(cookies={"access_token": "garbage"})) is None
|
|
|
|
|
|
def test_authenticate_websocket_query_token_rejected(monkeypatch):
|
|
"""BUG-036: the access token must never be accepted from the query string."""
|
|
monkeypatch.setenv("OBSIGATE_AUTH_ENABLED", "true")
|
|
from backend.auth.jwt_handler import create_access_token
|
|
|
|
token = create_access_token({
|
|
"username": "u", "role": "user", "vaults": ["*"], "display_name": "U",
|
|
})
|
|
ws = _StubWebSocket(query={"token": token})
|
|
assert authenticate_websocket(ws) is None
|
|
|
|
|
|
def test_authenticate_websocket_cookie_token_accepted(monkeypatch):
|
|
"""The HttpOnly access_token cookie remains the supported transport."""
|
|
monkeypatch.setenv("OBSIGATE_AUTH_ENABLED", "true")
|
|
import backend.auth.user_store as user_store
|
|
from backend.auth.jwt_handler import create_access_token
|
|
|
|
token = create_access_token({
|
|
"username": "u", "role": "user", "vaults": ["*"], "display_name": "U",
|
|
})
|
|
monkeypatch.setattr(user_store, "get_user", lambda username: {
|
|
"username": username, "role": "user", "vaults": ["*"],
|
|
"display_name": "U", "active": True,
|
|
})
|
|
user = authenticate_websocket(_StubWebSocket(cookies={"access_token": token}))
|
|
assert user is not None
|
|
assert user["username"] == "u"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Manager unit tests (no WebSocket transport)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_manager(tmp_path: Path, debounce: float = 0.05) -> CollabManager:
|
|
return CollabManager(save_debounce=debounce)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_persists_text_after_debounce(tmp_path: Path):
|
|
target = tmp_path / "note.md"
|
|
target.write_text("initial", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target)
|
|
|
|
room.pending_text = "updated content"
|
|
manager._schedule_save(room)
|
|
await asyncio.sleep(0.2)
|
|
|
|
assert target.read_text(encoding="utf-8") == "updated content"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_manager_flush_writes_pending_text(tmp_path: Path):
|
|
target = tmp_path / "note.md"
|
|
target.write_text("initial", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target, pending_text="flushed")
|
|
await manager._flush(room)
|
|
assert target.read_text(encoding="utf-8") == "flushed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_message_rejects_oversized_update(tmp_path: Path):
|
|
target = tmp_path / "note.md"
|
|
target.write_text("x", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target)
|
|
client = _FakeClient(conn_id=1)
|
|
|
|
payload = base64.b64encode(b"a" * 10).decode()
|
|
# Monkeypatch the max size check by sending an oversized base64 string.
|
|
import backend.collab as collab_mod
|
|
|
|
original = collab_mod.MAX_UPDATE_BYTES
|
|
try:
|
|
collab_mod.MAX_UPDATE_BYTES = 4
|
|
await manager._on_message(room, client, json.dumps({"type": "update", "update": payload}))
|
|
finally:
|
|
collab_mod.MAX_UPDATE_BYTES = original
|
|
|
|
assert room.updates == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_message_rejects_oversized_raw(tmp_path: Path):
|
|
"""BUG-036: oversized raw frames are dropped before parsing."""
|
|
target = tmp_path / "note.md"
|
|
target.write_text("x", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target)
|
|
client = _FakeClient(conn_id=1)
|
|
|
|
import backend.collab as collab_mod
|
|
|
|
original = collab_mod.MAX_MESSAGE_CHARS
|
|
try:
|
|
collab_mod.MAX_MESSAGE_CHARS = 10
|
|
await manager._on_message(room, client, json.dumps({"type": "text", "text": "hello"}))
|
|
finally:
|
|
collab_mod.MAX_MESSAGE_CHARS = original
|
|
|
|
assert room.pending_text is None
|
|
|
|
|
|
class _FakeWebSocket:
|
|
def __init__(self):
|
|
self.sent: list[dict] = []
|
|
|
|
async def send_json(self, message: dict):
|
|
self.sent.append(message)
|
|
|
|
|
|
class _FakeClient:
|
|
def __init__(self, conn_id: int):
|
|
self.conn_id = conn_id
|
|
self.websocket = _FakeWebSocket()
|
|
self.username = "u"
|
|
self.display_name = "u"
|
|
self.color = "#000"
|
|
self.y_client_id = None
|
|
self.awareness = None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_message_appends_update_and_broadcasts(tmp_path: Path):
|
|
target = tmp_path / "note.md"
|
|
target.write_text("x", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target)
|
|
sender = _FakeClient(1)
|
|
receiver = _FakeClient(2)
|
|
room.clients[1] = sender
|
|
room.clients[2] = receiver
|
|
|
|
encoded = base64.b64encode(b"hello-yjs").decode()
|
|
await manager._on_message(room, sender, json.dumps({"type": "update", "update": encoded}))
|
|
|
|
assert room.updates == [b"hello-yjs"]
|
|
assert room.has_updates is True
|
|
assert receiver.websocket.sent[-1]["type"] == "update"
|
|
assert receiver.websocket.sent[-1]["update"] == encoded
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_message_awareness_relayed(tmp_path: Path):
|
|
target = tmp_path / "note.md"
|
|
target.write_text("x", encoding="utf-8")
|
|
manager = _make_manager(tmp_path)
|
|
room = CollabRoom(vault="V", path="note.md", file_path=target)
|
|
sender = _FakeClient(1)
|
|
receiver = _FakeClient(2)
|
|
room.clients[1] = sender
|
|
room.clients[2] = receiver
|
|
|
|
msg = {"type": "awareness", "clientId": 42, "state": {"cursor": {"anchor": 1, "head": 2}}}
|
|
await manager._on_message(room, sender, json.dumps(msg))
|
|
|
|
assert sender.y_client_id == 42
|
|
assert receiver.websocket.sent[-1]["type"] == "awareness"
|
|
assert receiver.websocket.sent[-1]["state"]["cursor"]["head"] == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration tests (Starlette TestClient WebSocket)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_websocket_init_and_broadcast(client):
|
|
with client.websocket_connect("/ws/collab/TestVault/note1.md") as ws1:
|
|
init1 = _recv_json(ws1)
|
|
assert init1["type"] == "init"
|
|
assert init1["connId"] == 1
|
|
assert init1["seed"] is not None # first client receives the file content
|
|
assert init1["updates"] == []
|
|
|
|
with client.websocket_connect("/ws/collab/TestVault/note1.md") as ws2:
|
|
init2 = _recv_json(ws2)
|
|
assert init2["type"] == "init"
|
|
assert init2["seed"] is None # already seeded
|
|
# ws1 is notified of the new peer
|
|
joined = _recv_json(ws1)
|
|
assert joined["type"] == "peer_joined"
|
|
assert len(init2["peers"]) == 2
|
|
|
|
encoded = base64.b64encode(b"payload").decode()
|
|
ws1.send_text(json.dumps({"type": "update", "update": encoded}))
|
|
relayed = _recv_json(ws2)
|
|
assert relayed["type"] == "update"
|
|
assert relayed["update"] == encoded
|
|
|
|
# ws2 left → ws1 gets peer_left
|
|
left = _recv_json(ws1)
|
|
assert left["type"] == "peer_left"
|
|
|
|
|
|
def test_websocket_awareness_relay(client):
|
|
with client.websocket_connect("/ws/collab/TestVault/note1.md") as ws1:
|
|
_recv_json(ws1)
|
|
with client.websocket_connect("/ws/collab/TestVault/note1.md") as ws2:
|
|
_recv_json(ws2)
|
|
_recv_json(ws1) # peer_joined
|
|
|
|
ws1.send_text(json.dumps({
|
|
"type": "awareness",
|
|
"clientId": 7,
|
|
"state": {"user": {"name": "Alice"}, "cursor": {"anchor": 3, "head": 3}},
|
|
}))
|
|
msg = _recv_json(ws2)
|
|
assert msg["type"] == "awareness"
|
|
assert msg["clientId"] == 7
|
|
assert msg["state"]["user"]["name"] == "Alice"
|
|
|
|
|
|
def test_websocket_persists_text(client, test_vault_dir):
|
|
from backend import collab as collab_mod
|
|
|
|
original = collab_mod.collab_manager._save_debounce
|
|
collab_mod.collab_manager._save_debounce = 0.05
|
|
try:
|
|
with client.websocket_connect("/ws/collab/TestVault/note1.md") as ws:
|
|
_recv_json(ws)
|
|
ws.send_text(json.dumps({"type": "text", "text": "# Persisted by collab\n"}))
|
|
deadline = time.time() + 2
|
|
target = Path(test_vault_dir) / "note1.md"
|
|
while time.time() < deadline:
|
|
if target.read_text(encoding="utf-8") == "# Persisted by collab\n":
|
|
break
|
|
time.sleep(0.05)
|
|
assert target.read_text(encoding="utf-8") == "# Persisted by collab\n"
|
|
finally:
|
|
collab_mod.collab_manager._save_debounce = original
|
|
|
|
|
|
def test_websocket_unknown_file_rejected(client):
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect("/ws/collab/TestVault/missing.md") as ws:
|
|
ws.receive_text()
|
|
|
|
|
|
def test_websocket_unknown_vault_rejected(client):
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with client.websocket_connect("/ws/collab/Nope/note1.md") as ws:
|
|
ws.receive_text()
|
|
|
|
|
|
def test_websocket_requires_auth(admin_client):
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with admin_client.websocket_connect("/ws/collab/TestVault/note1.md") as ws:
|
|
ws.receive_text()
|
|
|
|
|
|
def test_websocket_vault_acl_enforced(admin_client):
|
|
from backend.indexer import get_vault_data
|
|
|
|
root = Path(get_vault_data("TestVault")["path"])
|
|
target = root / "acl_test.md"
|
|
target.write_text("hi", encoding="utf-8")
|
|
try:
|
|
_login(admin_client, "normaluser", "normal123") # sets the access_token cookie
|
|
# normaluser has access to TestVault only.
|
|
with admin_client.websocket_connect("/ws/collab/TestVault/acl_test.md") as ws:
|
|
assert _recv_json(ws)["type"] == "init"
|
|
|
|
with pytest.raises(WebSocketDisconnect):
|
|
with admin_client.websocket_connect("/ws/collab/OtherVault/acl_test.md") as ws:
|
|
ws.receive_text()
|
|
finally:
|
|
target.unlink(missing_ok=True)
|
|
|
|
|
|
def _login(client: TestClient, username: str, password: str) -> str:
|
|
resp = client.post("/api/auth/login", json={"username": username, "password": password})
|
|
assert resp.status_code == 200, resp.text
|
|
return resp.json()["access_token"]
|
|
|
|
|
|
def _recv_until(ws, msg_type: str, limit: int = 20) -> dict:
|
|
for _ in range(limit):
|
|
msg = _recv_json(ws)
|
|
if msg.get("type") == msg_type:
|
|
return msg
|
|
raise AssertionError(f"message type '{msg_type}' not received")
|
|
|
|
|
|
def test_websocket_load_five_clients(client):
|
|
"""5+ utilisateurs simultanés sur le même fichier (ROADMAP #62)."""
|
|
from backend.collab import collab_manager
|
|
|
|
contexts = []
|
|
try:
|
|
sessions = []
|
|
for _ in range(5):
|
|
ctx = client.websocket_connect("/ws/collab/TestVault/note1.md")
|
|
ws = ctx.__enter__()
|
|
contexts.append(ctx)
|
|
sessions.append(ws)
|
|
assert _recv_json(ws)["type"] == "init"
|
|
|
|
assert collab_manager.room_peer_count("TestVault", "note1.md") == 5
|
|
|
|
# The first client broadcasts an update, all 4 others receive it.
|
|
encoded = base64.b64encode(b"load-test-update").decode()
|
|
sessions[0].send_text(json.dumps({"type": "update", "update": encoded}))
|
|
for ws in sessions[1:]:
|
|
relayed = _recv_until(ws, "update")
|
|
assert relayed["update"] == encoded
|
|
finally:
|
|
for ctx in reversed(contexts):
|
|
try:
|
|
ctx.__exit__(None, None, None)
|
|
except Exception:
|
|
pass
|