Imago/tests/test_websocket.py

"""
Tests WebSocket — suivi pipeline temps réel.

Utilise le TestClient de Starlette pour les WebSockets.
"""
import pytest
from unittest.mock import patch

from starlette.testclient import TestClient

from app.main import app
from app.models.image import Image, ProcessingStatus
from app.models.client import APIClient


# Client synchrone pour les WebSockets (Starlette TestClient)
# On mocke complètement la DB pour éviter les soucis d'event loop (greenlet_spawn)
sync_client = TestClient(app)


# ─────────────────────────────────────────────────────────────
# Tests
# ─────────────────────────────────────────────────────────────

class TestWSPipelineNoToken:
    """Connexion sans token → refus immédiat."""

    def test_no_token_rejected(self):
        """WS sans token reçoit un close 4001."""
        with pytest.raises(Exception):
            with sync_client.websocket_connect("/ws/pipeline/1"):
                pass


@patch("app.routers.websocket._get_image")
@patch("app.routers.websocket._authenticate_ws")
class TestWSPipelineAuth:
    """Tests d'authentification WebSocket appliqués avec un mock de DB."""

    def test_invalid_token_rejected(self, mock_auth, mock_get_image):
        """WS avec un token invalide reçoit un close."""
        mock_auth.return_value = None
        with pytest.raises(Exception):
            with sync_client.websocket_connect("/ws/pipeline/1?token=invalid-key"):
                pass

    def test_valid_token_accepted(self, mock_auth, mock_get_image):
        """WS avec token valide d'un client qui possède l'image → accepté."""
        # Mock du client authentifié
        client_mock = APIClient(id="client-a", scopes=[])
        client_mock.has_scope = lambda x: False
        mock_auth.return_value = client_mock

        # Mock de l'image (DONE -> envoie message synthétique)
        image_mock = Image()
        image_mock.id = 1
        image_mock.client_id = "client-a"
        image_mock.processing_status = ProcessingStatus.DONE
        mock_get_image.return_value = image_mock

        with sync_client.websocket_connect(
            "/ws/pipeline/1?token=valid-key"
        ) as ws:
            data = ws.receive_json()
            assert data["event"] == "pipeline.done"
            assert data["synthetic"] is True

    def test_wrong_owner_rejected(self, mock_auth, mock_get_image):
        """WS avec token d'un client B sur image de client A → refus 4003."""
        # Mock du client authentifié (client B)
        client_mock = APIClient(id="client-b", scopes=[])
        client_mock.has_scope = lambda x: False
        mock_auth.return_value = client_mock

        # Mock de l'image appartenant au client A
        image_mock = Image()
        image_mock.id = 1
        image_mock.client_id = "client-a"
        image_mock.processing_status = ProcessingStatus.PENDING
        mock_get_image.return_value = image_mock

        with pytest.raises(Exception):
            with sync_client.websocket_connect(
                "/ws/pipeline/1?token=valid-key"
            ):
                pass


@patch("app.routers.websocket._get_image")
@patch("app.routers.websocket._authenticate_ws")
class TestWSPipelineSynthetic:
    """Tests des messages synthétiques."""

    def test_done_image_sends_synthetic(self, mock_auth, mock_get_image):
        """Image déjà 'done' → reçoit pipeline.done synthétique."""
        client_mock = APIClient(id="client-a", scopes=[])
        client_mock.has_scope = lambda x: False
        mock_auth.return_value = client_mock

        image_mock = Image()
        image_mock.id = 1
        image_mock.client_id = "client-a"
        image_mock.processing_status = ProcessingStatus.DONE
        mock_get_image.return_value = image_mock

        with sync_client.websocket_connect(
            "/ws/pipeline/1?token=valid-key"
        ) as ws:
            data = ws.receive_json()
            assert data["event"] == "pipeline.done"
            assert data["image_id"] == 1
            assert data["synthetic"] is True

    def test_error_image_sends_synthetic(self, mock_auth, mock_get_image):
        """Image déjà 'error' → reçoit pipeline.error synthétique."""
        client_mock = APIClient(id="client-a", scopes=[])
        client_mock.has_scope = lambda x: False
        mock_auth.return_value = client_mock

        image_mock = Image()
        image_mock.id = 1
        image_mock.client_id = "client-a"
        image_mock.processing_status = ProcessingStatus.ERROR
        image_mock.processing_error = "AI timeout"
        mock_get_image.return_value = image_mock

        with sync_client.websocket_connect(
            "/ws/pipeline/1?token=valid-key"
        ) as ws:
            data = ws.receive_json()
            assert data["event"] == "pipeline.error"
            assert data["image_id"] == 1
            assert "AI timeout" in data["error"]
            assert data["synthetic"] is True


@patch("app.routers.websocket._authenticate_ws")
class TestWSAdminMonitor:
    """Tests du monitoring admin."""

    def test_admin_no_token_rejected(self, mock_auth):
        """WS admin sans token → refus."""
        mock_auth.return_value = None
        with pytest.raises(Exception):
            with sync_client.websocket_connect("/ws/admin/monitor"):
                pass

    def test_admin_non_admin_rejected(self, mock_auth):
        """WS admin avec token non-admin → refus 4003."""
        client_mock = APIClient(id="client-a", scopes=[])
        client_mock.has_scope = lambda x: False
        mock_auth.return_value = client_mock

        with pytest.raises(Exception):
            with sync_client.websocket_connect(
                "/ws/admin/monitor?token=valid-key"
            ):
                pass