lct-hack/backend/tests/test_ws_ownership.py

"""Знание URL занятия не даёт курсанту доступ к чужому АРМ."""

import asyncio
import importlib
import json
import time
from uuid import uuid4

import pytest
from fastapi.testclient import TestClient

from app.api.auth import Principal
from app.api.ws import session as session_ws
from app.domain.roles import Role
from app.main import app
from app.scenarios import store
from app.session.hub import hub
from app.session.store import MemorySessionStore
from app.session.state import SessionState


@pytest.fixture
def client():
    with TestClient(app) as test_client:
        test_client.post("/api/auth/dev-token")
        hub.store = MemorySessionStore()
        yield test_client


def test_trainee_cannot_open_foreign_or_unassigned_call_or_station(client, monkeypatch):
    session_id = uuid4()
    with client.websocket_connect(f"/ws/control/{session_id}") as control:
        control.send_json({"type": "scenario.start", "scenario_id": "fire-apartment-l2",
                           "trainee": "Назначенный", "mode": "training"})
        deadline = time.monotonic() + 3
        while hub.get(session_id) is None and time.monotonic() < deadline:
            time.sleep(0.02)
        state = hub.get(session_id)
        assert state is not None

        owner_id, foreign_id = uuid4(), uuid4()
        modules = {
            "call": importlib.import_module("app.api.ws.call"),
            "station": importlib.import_module("app.api.ws.station"),
        }
        for assigned in (owner_id, None):
            state.trainee_id = assigned
            for role_name, module in modules.items():
                who = Principal(login="foreign", full_name="Чужой",
                                role=Role.TRAINEE, trainee_id=foreign_id)
                monkeypatch.setattr(session_ws, "principal_of", lambda _ws, user=who: user)
                with client.websocket_connect(f"/ws/{role_name}/{session_id}") as socket:
                    event = socket.receive_json()
                    assert event["type"] == "error" and event["code"] == "forbidden"

        state.trainee_id = owner_id
        for role_name, module in modules.items():
            who = Principal(login="owner", full_name="Назначенный",
                            role=Role.TRAINEE, trainee_id=owner_id)
            monkeypatch.setattr(session_ws, "principal_of", lambda _ws, user=who: user)
            with client.websocket_connect(f"/ws/{role_name}/{session_id}") as socket:
                socket.send_json({"type": "unknown"})
                for _ in range(8):
                    event = socket.receive_json()
                    if event["type"] == "error":
                        assert event["code"] == "unsupported_event"
                        break
                else:
                    raise AssertionError("владелец занятия не допущен")


def test_instructor_cannot_start_or_queue_another_instructors_scenario(client, monkeypatch):
    source = store.get("fire-apartment-l2")
    assert source is not None
    private = source.model_copy(update={"id": "private-ws-owner-test"}, deep=True)
    store.register_owned_scenario(private, "another-instructor")
    emitted = []
    monkeypatch.setattr(hub, "to_observers", lambda _session_id, event: emitted.append(event))
    public_id = "fire-apartment-l2"
    cases = [
        {"scenario_id": private.id},
        {"scenario_id": public_id, "exercise": "dds",
         "scenario_ids": [public_id, private.id]},
        {"scenario_id": public_id, "exercise": "dds",
         "random_scenario_ids": [public_id, private.id]},
    ]
    sessions = []
    try:
        for fields in cases:
            session_id = uuid4()
            sessions.append(session_id)
            with client.websocket_connect(f"/ws/control/{session_id}") as control:
                control.send_json({
                    "type": "scenario.start", "trainee": "Курсант", "mode": "training",
                    **fields,
                })
                deadline = time.monotonic() + 3
                while len(emitted) < len(sessions) and time.monotonic() < deadline:
                    time.sleep(0.02)
                assert emitted[-1].code.value == "scenario_invalid"
                assert hub.get(session_id) is None
    finally:
        for session_id in sessions:
            hub.drop(session_id)
        store._library.pop(private.id, None)
        store._demo_scenario_owners.pop(private.id, None)


def test_idle_control_socket_closes_when_backend_lease_becomes_uncertain(monkeypatch):
    control_module = importlib.import_module("app.api.ws.control")
    session_id = uuid4()
    state = SessionState(
        session_id=session_id, scenario_id="test", scenario_title="Тест", level="L1",
        mode="training",
    )
    hub.register(state)
    who = Principal(login="lease-owner", full_name="Преподаватель", role=Role.INSTRUCTOR)
    monkeypatch.setattr(session_ws, "websocket_origin_allowed", lambda _ws: True)
    monkeypatch.setattr(session_ws, "principal_of", lambda _ws: who)

    class IdleSocket:
        accepted = False
        closed = None
        sent = []

        async def accept(self):
            self.accepted = True

        async def receive_json(self):
            await asyncio.sleep(0.02)
            await hub.fence(state)
            await asyncio.sleep(5)

        async def send_text(self, message):
            self.sent.append(json.loads(message))

        async def close(self, code=None):
            self.closed = code

    socket = IdleSocket()
    try:
        asyncio.run(control_module.control(socket, session_id))
        assert socket.accepted
        assert socket.closed == 1012
        assert socket.sent[-1]["message"].endswith("переподключитесь.")
    finally:
        hub.drop(session_id)


def test_control_command_checkpoint_failure_returns_fencing_error(monkeypatch):
    control_module = importlib.import_module("app.api.ws.control")
    session_id = uuid4()
    state = SessionState(
        session_id=session_id, scenario_id="test", scenario_title="Тест", level="L1",
        mode="training", owner_login="lease-owner",
    )

    class BrokenStore(MemorySessionStore):
        async def commit(self, _state, _records=()):
            raise OSError("simulated database partition")

    class OneCommandSocket:
        accepted = False
        closed = None
        sent = []

        async def accept(self):
            self.accepted = True

        async def receive_json(self):
            return {"type": "reference.play"}

        async def send_text(self, message):
            self.sent.append(json.loads(message))

        async def close(self, code=None):
            self.closed = code

    who = Principal(login="lease-owner", full_name="Преподаватель", role=Role.INSTRUCTOR)
    socket = OneCommandSocket()
    old_store = hub.store
    hub.store = BrokenStore()
    hub.register(state)
    monkeypatch.setattr(session_ws, "websocket_origin_allowed", lambda _ws: True)
    monkeypatch.setattr(session_ws, "principal_of", lambda _ws: who)
    try:
        with hub.observer(session_id) as observer_queue:
            asyncio.run(control_module.control(socket, session_id))
            assert socket.accepted
            assert state.lease_fenced
            assert socket.closed == 1012
            assert socket.sent[-1]["type"] == "error"
            assert "переподключитесь" in socket.sent[-1]["message"]
            event = observer_queue.get_nowait()
            assert "переподключитесь" in event.message
            assert observer_queue.empty(), "uncommitted controller event leaked to observers"
    finally:
        hub.store = old_store
        hub.drop(session_id)


def test_control_does_not_lose_command_read_during_fencing_poll(monkeypatch):
    """Опрос fencing по таймеру не отменяет чтение, уже забравшее кадр из сокета.

    Сокет отдаёт команду, но возвращается из receive позже тика опроса — так
    выглядит задержка планировщика под нагрузкой. Отмена чтения на тике
    выбрасывает уже прочитанную команду.
    """
    control_module = importlib.import_module("app.api.ws.control")
    monkeypatch.setattr(control_module, "_FENCE_POLL_SECONDS", 0.02)
    session_id = uuid4()
    state = SessionState(
        session_id=session_id, scenario_id="test", scenario_title="Тест", level="L1",
        mode="training", owner_login="lease-owner",
    )
    hub.register(state)
    who = Principal(login="lease-owner", full_name="Преподаватель", role=Role.INSTRUCTOR)
    monkeypatch.setattr(session_ws, "websocket_origin_allowed", lambda _ws: True)
    monkeypatch.setattr(session_ws, "principal_of", lambda _ws: who)

    class SlowReturnSocket:
        def __init__(self):
            self.wire = [{"type": "no.such.command"}]
            self.lost = []
            self.closed = None
            self.sent = []

        async def accept(self):
            pass

        async def receive_json(self):
            if not self.wire:
                await hub.fence(state)
                await asyncio.sleep(5)
            message = self.wire.pop(0)
            try:
                await asyncio.sleep(0.1)
            except asyncio.CancelledError:
                self.lost.append(message)
                raise
            return message

        async def send_text(self, message):
            self.sent.append(json.loads(message))

        async def close(self, code=None):
            self.closed = code

    socket = SlowReturnSocket()
    try:
        with hub.observer(session_id) as observer_queue:
            asyncio.run(control_module.control(socket, session_id))
            assert socket.lost == []
            event = observer_queue.get_nowait()
            assert event.code == "unsupported_event"
            assert socket.closed == 1012
            assert socket.sent[-1]["message"].endswith("переподключитесь.")
    finally:
        hub.drop(session_id)