leocrm/tests/test_graceful_shutdown.py

"""Tests for graceful shutdown & connection draining (B.15).

Covers:
- SIGTERM → shutdown event set
- WS drain → all connections closed
- Worker shutdown → on_shutdown called
"""

from __future__ import annotations

import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from starlette.websockets import WebSocket

from app.core.ws_helpers import (
    drain_all_connections,
    register_ws_registry,
    _global_ws_registries,
)


class TestShutdownEvent:
    """Tests for the API graceful shutdown event."""

    def test_shutdown_event_is_settable(self):
        """The _shutdown_event should be settable."""
        from app.main import _shutdown_event
        # Reset to clean state
        _shutdown_event.clear()
        assert not _shutdown_event.is_set()
        _shutdown_event.set()
        assert _shutdown_event.is_set()
        # Clean up
        _shutdown_event.clear()

    def test_shutdown_event_is_asyncio_event(self):
        """_shutdown_event should be an asyncio.Event."""
        from app.main import _shutdown_event
        assert isinstance(_shutdown_event, asyncio.Event)


class TestWebSocketDrain:
    """Tests for WebSocket connection draining."""

    @pytest.mark.asyncio
    async def test_drain_all_connections_closes_websockets(self):
        """drain_all_connections should close all registered WS connections."""
        # Create mock WebSocket objects
        ws1 = MagicMock(spec=WebSocket)
        ws1.send_text = AsyncMock()
        ws1.close = AsyncMock()
        ws2 = MagicMock(spec=WebSocket)
        ws2.send_text = AsyncMock()
        ws2.close = AsyncMock()

        # Register a test registry
        test_registry: dict[str, list[WebSocket]] = {
            "user1": [ws1],
            "user2": [ws2],
        }
        register_ws_registry(test_registry)

        # Drain with 0 grace period for fast test
        await drain_all_connections(grace_period_seconds=0)

        # Both WebSockets should have received reconnect message and been closed
        ws1.send_text.assert_called_once()
        ws2.send_text.assert_called_once()
        ws1.close.assert_called_once()
        ws2.close.assert_called_once()

        # Verify reconnect message content
        call_args = ws1.send_text.call_args[0][0]
        msg = json.loads(call_args)
        assert msg["type"] == "reconnect"
        assert msg["reason"] == "server_shutdown"

    @pytest.mark.asyncio
    async def test_drain_all_connections_handles_empty_registry(self):
        """drain_all_connections should handle empty registries gracefully."""
        # Clear any existing registries from previous tests
        _global_ws_registries.clear()
        await drain_all_connections(grace_period_seconds=0)
        # Should not raise

    @pytest.mark.asyncio
    async def test_drain_all_connections_handles_send_failure(self):
        """drain_all_connections should not fail if send_text raises."""
        ws = MagicMock(spec=WebSocket)
        ws.send_text = AsyncMock(side_effect=Exception("WS closed"))
        ws.close = AsyncMock(side_effect=Exception("WS closed"))

        test_registry: dict[str, list[WebSocket]] = {"user1": [ws]}
        register_ws_registry(test_registry)

        # Should not raise despite send failures
        await drain_all_connections(grace_period_seconds=0)

    @pytest.mark.asyncio
    async def test_register_ws_registry_dedup(self):
        """register_ws_registry should not add the same registry twice."""
        _global_ws_registries.clear()
        reg: dict[str, list[WebSocket]] = {}
        register_ws_registry(reg)
        register_ws_registry(reg)
        assert len(_global_ws_registries) == 1


class TestWorkerShutdown:
    """Tests for ARQ worker graceful shutdown."""

    @pytest.mark.asyncio
    async def test_on_shutdown_called(self):
        """on_shutdown should be callable and close Redis."""
        from app.core.worker import on_shutdown

        with patch("app.core.auth.close_redis", new_callable=AsyncMock) as mock_close:
            with patch("app.core.db.get_worker_session_factory") as mock_factory:
                # Mock the session factory to avoid DB access
                mock_session = AsyncMock()
                mock_result = MagicMock()
                mock_result.scalars.return_value.all.return_value = []
                mock_session.execute = AsyncMock(return_value=mock_result)
                mock_session.commit = AsyncMock()
                mock_ctx = AsyncMock()
                mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
                mock_ctx.__aexit__ = AsyncMock(return_value=False)
                mock_factory.return_value = mock_ctx

                await on_shutdown({})
                mock_close.assert_called_once()

    @pytest.mark.asyncio
    async def test_on_shutdown_handles_db_error(self):
        """on_shutdown should handle DB errors gracefully and still close Redis."""
        from app.core.worker import on_shutdown

        with patch("app.core.auth.close_redis", new_callable=AsyncMock) as mock_close:
            with patch("app.core.db.get_worker_session_factory", side_effect=Exception("DB unavailable")):
                await on_shutdown({})
                mock_close.assert_called_once()