tests/integration/test_client_disconnect_advanced.py

"""
Advanced integration tests for client disconnect detection during request processing.

These tests use REAL asyncio primitives to verify disconnect detection works
correctly in real-world async scenarios with actual timing and concurrency.

Test Strategy:
- Use real_server_state fixture (real asyncio.Lock, asyncio.Queue, asyncio.Event)
- Mock only external boundaries (browser, HTTP request)
- Test disconnect detection at various stages of processing
- Verify proper cleanup and resource release

Coverage Target: Disconnect detection integrity and resource cleanup
"""

import asyncio
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import HTTPException
from fastapi.responses import JSONResponse

from api_utils.context_types import QueueItem


# Test stub for QueueManager used in these tests
class QueueManager:
    """Functional stub class for queue management in tests."""

    request_queue: Any = None
    processing_lock: Any = None
    logger: Any = None

    async def process_request(self, item: QueueItem) -> Any:
        """Call real process_request with item details."""
        from api_utils.request_processor import process_request

        return await process_request(
            item["req_id"],
            item["request_data"],
            item["http_request"],
            item["result_future"],
        )

    async def check_queue_disconnects(self) -> Any:
        """Stub check_queue_disconnects method."""
        from api_utils.client_connection import check_client_connection
        from api_utils.error_utils import client_disconnected

        queue_size = self.request_queue.qsize()
        items_to_requeue = []
        for _ in range(queue_size):
            try:
                item = self.request_queue.get_nowait()
                item_req_id = item.get("req_id")
                if not item.get("cancelled", False):
                    item_http_req = item.get("http_request")
                    if item_http_req:
                        try:
                            if not await check_client_connection(
                                item_req_id, item_http_req
                            ):
                                item["cancelled"] = True
                                item_fut = item.get("result_future")
                                if item_fut and not item_fut.done():
                                    item_fut.set_exception(
                                        client_disconnected(
                                            item_req_id,
                                            "Client disconnected while queued.",
                                        )
                                    )
                        except Exception as e:
                            self.logger.error(
                                f"[{item_req_id}] Error in disconnect check: {e}"
                            )
                items_to_requeue.append(item)
            except asyncio.QueueEmpty:
                break
        for item in items_to_requeue:
            await self.request_queue.put(item)


@pytest.mark.integration
class TestClientDisconnectDuringQueueWait:
    """Tests for disconnect detection while request waits in queue."""

    async def test_disconnect_detected_before_processing_starts(
        self, real_server_state
    ):
        """
        Test client disconnect is detected before request starts processing.

        Uses real asyncio.Lock to simulate lock contention scenario.
        """
        queue = real_server_state.request_queue
        lock = real_server_state.processing_lock
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.processing_lock = lock
        queue_manager.logger = MagicMock()

        # Mock task_done since we call process_request directly without queue.get()
        queue.task_done = MagicMock()

        # Create first request that holds the lock
        async def hold_lock_request():
            async with lock:
                # Hold lock for 0.5 seconds
                await asyncio.sleep(0.5)

        # Start first request (holds lock)
        lock_holder = asyncio.create_task(hold_lock_request())

        try:
            # Wait to ensure lock is acquired
            await asyncio.sleep(0.05)

            # Create second request that will wait for lock
            req_item = cast(
                QueueItem,
                {
                    "req_id": "waiting-req",
                    "request_data": MagicMock(),
                    "http_request": MagicMock(),
                    "result_future": asyncio.Future(),
                    "cancelled": False,
                    "enqueue_time": 0.0,
                },
            )

            # Client disconnects immediately
            req_item["http_request"].is_disconnected = AsyncMock(return_value=True)

            # Try to process (should detect disconnect before acquiring lock)
            with patch(
                "api_utils.request_processor._check_client_connection",
                new_callable=AsyncMock,
                return_value=False,
            ):
                await queue_manager.process_request(req_item)

            # Verify disconnect was detected and 499 error set
            assert req_item["result_future"].done()
            with pytest.raises(HTTPException) as exc:
                req_item["result_future"].result()
            assert exc.value.status_code == 499

            # Clear the exception from the future to prevent warning
            try:
                req_item["result_future"].exception()
            except Exception:
                pass

        finally:
            # Cleanup: cancel and wait for lock holder to finish
            lock_holder.cancel()
            try:
                await lock_holder
            except asyncio.CancelledError:
                pass

    async def test_multiple_disconnects_in_queue(self, real_server_state):
        """
        Test multiple client disconnects while waiting in queue.

        Verifies check_queue_disconnects works with real asyncio.Queue.
        """
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        # Create 5 requests: 1st, 3rd, 5th are disconnected
        disconnect_map = {0: True, 1: False, 2: True, 3: False, 4: True}
        items = []

        for i in range(5):
            item = cast(
                QueueItem,
                {
                    "req_id": f"req-{i}",
                    "http_request": MagicMock(),
                    "cancelled": False,
                    "result_future": asyncio.Future(),
                    "request_data": None,
                    "enqueue_time": 0.0,
                },
            )
            item["http_request"].is_disconnected = AsyncMock(
                return_value=disconnect_map[i]
            )
            items.append(item)
            await queue.put(item)

        # Check disconnects
        await queue_manager.check_queue_disconnects()

        # Verify disconnected items were cancelled
        assert items[0]["cancelled"] is True
        assert items[1]["cancelled"] is False
        assert items[2]["cancelled"] is True
        assert items[3]["cancelled"] is False
        assert items[4]["cancelled"] is True

        # Verify all futures for disconnected items are done with 499
        for i in [0, 2, 4]:
            assert items[i]["result_future"].done()
            with pytest.raises(HTTPException) as exc:
                items[i]["result_future"].result()
            assert exc.value.status_code == 499

        # Clear exceptions from futures to prevent warnings
        for i in [0, 2, 4]:
            try:
                items[i]["result_future"].exception()
            except Exception:
                pass


@pytest.mark.integration
class TestClientDisconnectDuringProcessing:
    """Tests for disconnect detection during active request processing."""

    async def test_disconnect_during_lock_acquisition(self, real_server_state):
        """
        Test disconnect detected while waiting to acquire processing_lock.

        Simulates scenario where client disconnects during lock contention.
        """
        queue = real_server_state.request_queue
        lock = real_server_state.processing_lock
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.processing_lock = lock
        queue_manager.logger = MagicMock()

        # Mock task_done since we call process_request directly without queue.get()
        queue.task_done = MagicMock()

        disconnect_tracker = {"value": False}

        # Hold lock with first task
        async def lock_holder():
            async with lock:
                await asyncio.sleep(0.3)

        lock_task = asyncio.create_task(lock_holder())
        await asyncio.sleep(0.05)  # Ensure lock is held

        # Create request that will disconnect while waiting for lock
        req_item = cast(
            QueueItem,
            {
                "req_id": "disconnect-req",
                "request_data": MagicMock(),
                "http_request": MagicMock(),
                "result_future": asyncio.Future(),
                "cancelled": False,
                "enqueue_time": 0.0,
            },
        )

        # Simulate delayed disconnect (after initial check but before lock acquired)
        async def disconnect_checker(*args):
            """First check passes, second check (inside lock) fails."""
            if disconnect_tracker["value"]:
                return False  # Disconnected
            else:
                disconnect_tracker["value"] = True
                return True  # Still connected

        with patch(
            "api_utils.request_processor._check_client_connection",
            new_callable=AsyncMock,
            side_effect=disconnect_checker,
        ):
            await queue_manager.process_request(req_item)

        # Should detect disconnect and return 499
        assert req_item["result_future"].done()
        with pytest.raises(HTTPException) as exc:
            req_item["result_future"].result()
        assert exc.value.status_code == 499

        # Clear exception to prevent warning
        try:
            req_item["result_future"].exception()
        except Exception:
            pass

        await lock_task

    async def test_disconnect_during_streaming_response(self, real_server_state):
        """
        Test disconnect detection during streaming response generation.

        Simplified test: verify disconnect is detected and 499 error is set.
        """
        queue = real_server_state.request_queue
        lock = real_server_state.processing_lock
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.processing_lock = lock
        queue_manager.logger = MagicMock()

        # Mock task_done since we call process_request directly without queue.get()
        queue.task_done = MagicMock()

        req_item = cast(
            QueueItem,
            {
                "req_id": "streaming-req",
                "request_data": MagicMock(stream=True),
                "http_request": MagicMock(),
                "result_future": asyncio.Future(),
                "cancelled": False,
                "enqueue_time": 0.0,
            },
        )

        # Simulate disconnect after initial connection check
        disconnect_calls = [False, True]  # First call passes, second fails

        async def check_disconnect(*args):
            if disconnect_calls:
                return not disconnect_calls.pop(0)  # Invert for disconnect detection
            return True

        req_item["http_request"].is_disconnected = AsyncMock(
            side_effect=check_disconnect
        )

        # Mock check connection to trigger disconnect path
        with patch(
            "api_utils.request_processor._check_client_connection",
            new_callable=AsyncMock,
            side_effect=[True, False],  # Connected then disconnected
        ):
            await queue_manager.process_request(req_item)

        # Verify disconnect error was set
        assert req_item["result_future"].done()
        with pytest.raises(HTTPException) as exc:
            req_item["result_future"].result()
        assert exc.value.status_code == 499

        # Clean up future exception
        try:
            req_item["result_future"].exception()
        except Exception:
            pass


@pytest.mark.integration
class TestClientDisconnectRaceConditions:
    """Tests for race conditions between disconnect detection and processing."""

    async def test_concurrent_disconnect_checks_no_race(self, real_server_state):
        """
        Test that concurrent disconnect checks don't race.

        Multiple requests checking disconnect status simultaneously.
        """
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        # Create 10 concurrent requests
        items = []
        for i in range(10):
            item = cast(
                QueueItem,
                {
                    "req_id": f"concurrent-{i}",
                    "http_request": MagicMock(),
                    "cancelled": False,
                    "result_future": asyncio.Future(),
                    "request_data": None,
                    "enqueue_time": 0.0,
                },
            )
            # Half disconnected, half connected
            item["http_request"].is_disconnected = AsyncMock(return_value=i % 2 == 0)
            items.append(item)
            await queue.put(item)

        # Check disconnects (concurrent operations)
        await queue_manager.check_queue_disconnects()

        # Verify correct items marked as cancelled
        for i in range(10):
            if i % 2 == 0:
                assert items[i]["cancelled"] is True
                assert items[i]["result_future"].done()
            else:
                assert items[i]["cancelled"] is False

    async def test_disconnect_after_future_already_set(self, real_server_state):
        """
        Test disconnect check when result_future is already done.

        Verifies no race condition when future is set by processing
        before disconnect check happens.
        """
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        item = cast(
            QueueItem,
            {
                "req_id": "already-done",
                "http_request": MagicMock(),
                "cancelled": False,
                "result_future": asyncio.Future(),
                "request_data": None,
                "enqueue_time": 0.0,
            },
        )

        # Set future BEFORE disconnect check
        item["result_future"].set_result(
            JSONResponse(content={"message": "Already completed"})
        )

        # Now try to mark as disconnected
        item["http_request"].is_disconnected = AsyncMock(return_value=True)

        await queue.put(item)
        await queue_manager.check_queue_disconnects()

        # Should still be marked cancelled, but future should have original result
        assert item["cancelled"] is True
        result = item["result_future"].result()
        assert isinstance(result, JSONResponse)
        import json

        assert json.loads(bytes(result.body).decode()) == {
            "message": "Already completed"
        }


@pytest.mark.integration
class TestClientDisconnectCleanup:
    """Tests for proper cleanup after client disconnect."""

    async def test_cleanup_called_even_on_disconnect(self, real_server_state):
        """
        Test that resources are released when client disconnects.

        Simplified test: verify lock is released and future has error.
        """
        queue = real_server_state.request_queue
        lock = real_server_state.processing_lock
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.processing_lock = lock
        queue_manager.logger = MagicMock()

        # Mock task_done since we call process_request directly without queue.get()
        queue.task_done = MagicMock()

        req_item = cast(
            QueueItem,
            {
                "req_id": "cleanup-test",
                "request_data": MagicMock(),
                "http_request": MagicMock(),
                "result_future": asyncio.Future(),
                "cancelled": False,
                "enqueue_time": 0.0,
            },
        )

        # Mock check connection to trigger disconnect path
        with patch(
            "api_utils.request_processor._check_client_connection",
            new_callable=AsyncMock,
            side_effect=[True, False],  # Connected then disconnected
        ):
            await queue_manager.process_request(req_item)

        # Verify disconnect was handled
        assert req_item["result_future"].done()
        with pytest.raises(HTTPException) as exc:
            req_item["result_future"].result()
        assert exc.value.status_code == 499

        # Verify lock was released (cleanup successful)
        assert not lock.locked()

        # Clean up future exception
        try:
            req_item["result_future"].exception()
        except Exception:
            pass

    async def test_lock_released_on_disconnect(self, real_server_state):
        """
        Test that processing_lock is released when client disconnects.

        Critical for preventing deadlocks when clients disconnect.
        """
        lock = real_server_state.processing_lock
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.processing_lock = lock
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        # Mock task_done since we call process_request directly without queue.get()
        queue.task_done = MagicMock()

        req_item = cast(
            QueueItem,
            {
                "req_id": "lock-release-test",
                "request_data": MagicMock(),
                "http_request": MagicMock(),
                "result_future": asyncio.Future(),
                "cancelled": False,
                "enqueue_time": 0.0,
            },
        )

        # Simulate disconnect inside lock
        with patch(
            "api_utils.request_processor._check_client_connection",
            new_callable=AsyncMock,
            side_effect=[True, False],  # Disconnect inside lock
        ):
            await queue_manager.process_request(req_item)

        # Verify lock is released (can be acquired immediately)
        lock_acquired = lock.locked()
        assert not lock_acquired  # Lock should be released

        # Verify we can acquire it
        async with lock:
            assert lock.locked()
        assert not lock.locked()  # Released after context exit

        # Clean up future if it has exception
        if req_item["result_future"].done():
            try:
                req_item["result_future"].exception()
            except Exception:
                pass


@pytest.mark.integration
class TestClientDisconnectEdgeCases:
    """Tests for edge cases in disconnect detection."""

    async def test_disconnect_check_exception_handled(self, real_server_state):
        """
        Test that exceptions during disconnect check are handled gracefully.

        Network errors during is_disconnected() should not crash queue.
        """
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        # Create item with failing disconnect check
        item = cast(
            QueueItem,
            {
                "req_id": "error-check",
                "http_request": MagicMock(),
                "cancelled": False,
                "result_future": asyncio.Future(),
                "request_data": None,
                "enqueue_time": 0.0,
            },
        )
        item["http_request"].is_disconnected = AsyncMock(
            side_effect=Exception("Connection check failed")
        )

        await queue.put(item)

        # Should not raise, just log error
        await queue_manager.check_queue_disconnects()

        # Item should still be in queue
        assert queue.qsize() == 1

        # Error should be logged
        queue_manager.logger.error.assert_called()

    async def test_rapid_disconnect_reconnect_scenario(self, real_server_state):
        """
        Test rapid disconnect/reconnect scenario.

        Client disconnects then reconnects quickly - last state wins.
        """
        queue = real_server_state.request_queue
        queue_manager = QueueManager()
        queue_manager.request_queue = queue
        queue_manager.logger = MagicMock()

        disconnect_state = {"calls": 0}

        # Simulate alternating disconnect states
        async def alternating_disconnect(*args):
            disconnect_state["calls"] += 1
            # Odd calls = disconnected, even = connected
            return disconnect_state["calls"] % 2 == 1

        item = cast(
            QueueItem,
            {
                "req_id": "rapid-disconnect",
                "http_request": MagicMock(),
                "cancelled": False,
                "result_future": asyncio.Future(),
                "request_data": None,
                "enqueue_time": 0.0,
            },
        )
        item["http_request"].is_disconnected = AsyncMock(
            side_effect=alternating_disconnect
        )

        await queue.put(item)

        # First check: disconnected
        await queue_manager.check_queue_disconnects()
        assert item["cancelled"] is True

        # Put back and check again: connected (but already cancelled)
        await queue.put(item)
        await queue_manager.check_queue_disconnects()

        # Should remain cancelled (can't un-cancel)
        assert item["cancelled"] is True