ObsiGate/tests/test_collab.py

# 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


# ---------------------------------------------------------------------------
# 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 == []


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