tests/api_utils/routers/test_logs_ws.py

from unittest.mock import AsyncMock, MagicMock

import pytest
from fastapi import WebSocket, WebSocketDisconnect

from api_utils.routers.logs_ws import websocket_log_endpoint


@pytest.mark.asyncio
async def test_websocket_log_endpoint_success():
    """Test successful websocket connection and message loop."""
    websocket = AsyncMock(spec=WebSocket)
    logger = MagicMock()
    manager = AsyncMock()

    # Simulate one message then disconnect
    websocket.receive_text.side_effect = ["ping", WebSocketDisconnect()]

    await websocket_log_endpoint(websocket, logger, manager)

    # Verify connection
    manager.connect.assert_called_once()
    assert websocket.receive_text.call_count == 2

    # Verify disconnect
    manager.disconnect.assert_called_once()


@pytest.mark.asyncio
async def test_websocket_log_endpoint_no_manager():
    """Test endpoint when manager is None."""
    websocket = AsyncMock(spec=WebSocket)
    logger = MagicMock()

    await websocket_log_endpoint(websocket, logger, None)  # type: ignore[arg-type]

    websocket.close.assert_called_with(code=1011)


@pytest.mark.asyncio
async def test_websocket_log_endpoint_exception():
    """Test endpoint handling generic exception."""
    websocket = AsyncMock(spec=WebSocket)
    logger = MagicMock()
    manager = AsyncMock()

    # Simulate exception during receive
    websocket.receive_text.side_effect = Exception("Unexpected error")

    await websocket_log_endpoint(websocket, logger, manager)

    # Verify error logging
    logger.error.assert_called_once()

    # Verify disconnect is still called
    manager.disconnect.assert_called_once()