tests/integration/services/test_websocket_manager_coverage.py

"""
WebSocket Manager Coverage Tests

Comprehensive tests for WebSocketConnectionManager and DebuggingWebSocketManager
using AsyncMock patterns proven in Phase 169.
Target: 70%+ line coverage on core/websocket_manager.py (473 lines)

Tests use AsyncMock for all WebSocket async operations while testing real
connection tracking, broadcast logic, and debugging features.
"""

import pytest
from unittest.mock import AsyncMock, Mock, patch
from fastapi import WebSocket
from datetime import datetime

from core.websocket_manager import (
    WebSocketConnectionManager,
    DebuggingWebSocketManager,
    get_websocket_manager,
    get_debugging_websocket_manager,
)

# ============================================================================
# Fixtures
# ============================================================================

@pytest.fixture
def mock_websocket():
    """Create mock WebSocket with AsyncMock for async methods."""
    mock_ws = Mock(spec=WebSocket)
    mock_ws.accept = AsyncMock()
    mock_ws.send_text = AsyncMock()
    mock_ws.send_json = AsyncMock()
    mock_ws.close = AsyncMock()
    return mock_ws


@pytest.fixture
def manager():
    """Create fresh WebSocket manager for each test."""
    return WebSocketConnectionManager()


@pytest.fixture
def debug_manager(manager):
    """Create debugging manager with base manager."""
    return DebuggingWebSocketManager(manager)

# ============================================================================
# Connection Lifecycle Tests
# ============================================================================

class TestWebSocketConnectionLifecycle:
    """Test WebSocket connection lifecycle management."""

    @pytest.mark.asyncio
    async def test_connect_and_track_connection(self, manager, mock_websocket):
        """Test that WebSocket connections are tracked properly after connect."""
        await manager.connect(mock_websocket, stream_id="test_stream")

        # Verify connection was accepted
        mock_websocket.accept.assert_called_once()

        # Verify connection is tracked
        assert "test_stream" in manager.active_connections
        assert mock_websocket in manager.active_connections["test_stream"]

        # Verify connection metadata
        assert mock_websocket in manager.connection_streams
        assert manager.connection_streams[mock_websocket] == "test_stream"
        assert "connected_at" in manager.connection_info[mock_websocket]

    @pytest.mark.asyncio
    async def test_connect_sends_welcome_message(self, manager, mock_websocket):
        """Test that connect sends a welcome message to the client."""
        await manager.connect(mock_websocket, stream_id="test_stream")

        # Verify welcome message sent
        mock_websocket.send_text.assert_called()
        call_args = mock_websocket.send_text.call_args
        import json
        message = json.loads(call_args[0][0])
        assert message["type"] == "connected"
        assert message["stream_id"] == "test_stream"
        assert "timestamp" in message

    @pytest.mark.asyncio
    async def test_disconnect_removes_from_streams(self, manager, mock_websocket):
        """Test that disconnecting removes WebSocket from all tracking structures."""
        await manager.connect(mock_websocket, stream_id="stream1")
        manager.disconnect(mock_websocket)

        # Verify removal from active connections
        assert mock_websocket not in manager.active_connections.get("stream1", set())

        # Verify removal from connection streams
        assert mock_websocket not in manager.connection_streams

        # Verify removal from connection info
        assert mock_websocket not in manager.connection_info

    @pytest.mark.asyncio
    async def test_multiple_connections_per_stream(self, manager):
        """Test that multiple WebSockets can subscribe to the same stream."""
        mock_ws1 = Mock(spec=WebSocket)
        mock_ws2 = Mock(spec=WebSocket)
        mock_ws1.accept = AsyncMock()
        mock_ws2.accept = AsyncMock()
        mock_ws1.send_text = AsyncMock()
        mock_ws2.send_text = AsyncMock()

        await manager.connect(mock_ws1, stream_id="shared_stream")
        await manager.connect(mock_ws2, stream_id="shared_stream")

        # Both connections tracked
        assert len(manager.active_connections["shared_stream"]) == 2
        assert mock_ws1 in manager.active_connections["shared_stream"]
        assert mock_ws2 in manager.active_connections["shared_stream"]

    @pytest.mark.asyncio
    async def test_connection_cleanup_on_empty_stream(self, manager, mock_websocket):
        """Test that empty streams are removed from active_connections."""
        await manager.connect(mock_websocket, stream_id="temp_stream")
        assert "temp_stream" in manager.active_connections

        # Disconnect the only connection
        manager.disconnect(mock_websocket)

        # Stream should be removed
        assert "temp_stream" not in manager.active_connections

# ============================================================================
# Broadcast Tests
# ============================================================================

class TestWebSocketBroadcast:
    """Test WebSocket broadcasting functionality."""

    @pytest.mark.asyncio
    async def test_broadcast_to_all_subscribers(self, manager):
        """Test that broadcast reaches all subscribers in a stream."""
        mock_ws1 = Mock(spec=WebSocket)
        mock_ws2 = Mock(spec=WebSocket)
        mock_ws1.accept = AsyncMock()
        mock_ws2.accept = AsyncMock()
        mock_ws1.send_text = AsyncMock()
        mock_ws2.send_text = AsyncMock()

        await manager.connect(mock_ws1, stream_id="stream1")
        await manager.connect(mock_ws2, stream_id="stream1")

        message = {"type": "test", "data": "hello"}
        sent_count = await manager.broadcast("stream1", message)

        # Verify both received message
        assert sent_count == 2
        # send_text is called twice: once for welcome message, once for broadcast
        assert mock_ws1.send_text.call_count == 2
        assert mock_ws2.send_text.call_count == 2

        # Verify message format (last call is the broadcast)
        import json
        call_args = mock_ws1.send_text.call_args[0][0]
        sent_message = json.loads(call_args)
        assert sent_message["type"] == "test"
        assert sent_message["data"] == "hello"

    @pytest.mark.asyncio
    async def test_broadcast_handles_failed_connection(self, manager):
        """Test that broadcast handles connection failures gracefully."""
        mock_ws1 = Mock(spec=WebSocket)
        mock_ws2 = Mock(spec=WebSocket)
        mock_ws1.accept = AsyncMock()
        mock_ws2.accept = AsyncMock()
        mock_ws1.send_text = AsyncMock()
        mock_ws2.send_text = AsyncMock(side_effect=Exception("Connection lost"))
        mock_ws1.send_json = AsyncMock()
        mock_ws2.send_json = AsyncMock()

        await manager.connect(mock_ws1, stream_id="stream1")
        await manager.connect(mock_ws2, stream_id="stream1")

        # Broadcast (one will fail)
        sent_count = await manager.broadcast("stream1", {"type": "test"})

        # Should return successful sends only
        assert sent_count == 1

        # Failed connection should be disconnected
        assert mock_ws2 not in manager.active_connections.get("stream1", set())

    @pytest.mark.asyncio
    async def test_broadcast_to_empty_stream(self, manager):
        """Test broadcast to stream with no subscribers."""
        sent_count = await manager.broadcast("empty_stream", {"type": "test"})

        assert sent_count == 0

    @pytest.mark.asyncio
    async def test_send_personal_succeeds(self, manager, mock_websocket):
        """Test sending a message to a specific connection."""
        mock_websocket.send_text = AsyncMock(return_value=True)

        success = await manager.send_personal(
            mock_websocket,
            {"type": "personal", "data": "test"}
        )

        assert success is True
        mock_websocket.send_text.assert_called_once()

    @pytest.mark.asyncio
    async def test_send_personal_handles_failure(self, manager, mock_websocket):
        """Test send_personal handles send failure."""
        mock_websocket.send_text = AsyncMock(side_effect=Exception("Send failed"))

        success = await manager.send_personal(
            mock_websocket,
            {"type": "personal", "data": "test"}
        )

        assert success is False
        # WebSocket should be disconnected after failed send
        assert mock_websocket not in manager.connection_streams

    @pytest.mark.asyncio
    async def test_broadcast_trace_update(self, manager):
        """Test trace update broadcasting."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await manager.connect(mock_ws, stream_id="stream1")

        sent_count = await manager.broadcast_trace_update(
            "stream1",
            {"step": 1, "status": "running"}
        )

        assert sent_count == 1
        # Verify message structure
        call_args = mock_ws.send_text.call_args[0][0]
        import json
        message = json.loads(call_args)
        assert message["type"] == "trace_update"
        assert message["data"]["step"] == 1

    @pytest.mark.asyncio
    async def test_broadcast_session_update(self, manager):
        """Test debug session update broadcasting."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await manager.broadcast_session_update(
            "test123",
            "variable_changed",
            {"var": "value"}
        )

        assert sent_count == 1
        # Verify message routed to correct stream
        assert "debug_session_test123" in manager.active_connections

    @pytest.mark.asyncio
    async def test_broadcast_to_workspace(self, manager):
        """Test workspace broadcast."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await manager.connect(mock_ws, stream_id="workspace_test_ws")

        sent_count = await manager.broadcast_to_workspace(
            "test_ws",
            {"type": "workspace_update"}
        )

        assert sent_count == 1

# ============================================================================
# Stream Info Tests
# ============================================================================

class TestWebSocketStreamInfo:
    """Test stream information and statistics."""

    def test_get_connection_count(self, manager):
        """Test getting connection count for a stream."""
        assert manager.get_connection_count("non_existent") == 0

        # Add mock connection
        manager.active_connections["stream1"] = {Mock(), Mock()}
        assert manager.get_connection_count("stream1") == 2

    def test_get_all_streams(self, manager):
        """Test getting all active stream IDs."""
        manager.active_connections = {
            "stream1": {Mock()},
            "stream2": {Mock(), Mock()}
        }

        streams = manager.get_all_streams()
        assert "stream1" in streams
        assert "stream2" in streams
        assert len(streams) == 2

    @pytest.mark.asyncio
    async def test_get_stream_info(self, manager, mock_websocket):
        """Test getting detailed information about a stream."""
        await manager.connect(mock_websocket, stream_id="test_stream")

        info = manager.get_stream_info("test_stream")

        assert info["stream_id"] == "test_stream"
        assert info["connection_count"] == 1
        assert "connected_at" in info["connections"][0]

    def test_get_stream_info_for_non_existent(self, manager):
        """Test get_stream_info for non-existent stream."""
        info = manager.get_stream_info("non_existent")

        assert info["stream_id"] == "non_existent"
        assert info["connection_count"] == 0
        assert info["connections"] == []

# ============================================================================
# DebuggingWebSocketManager Tests
# ============================================================================

class TestDebuggingWebSocketManager:
    """Test debugging-specific WebSocket manager functionality."""

    @pytest.mark.asyncio
    async def test_stream_trace(self, debug_manager):
        """Test trace streaming."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="trace_exec1_session1")

        sent_count = await debug_manager.stream_trace(
            execution_id="exec1",
            session_id="session1",
            trace_data={"step": 1, "status": "running"}
        )

        assert sent_count == 1
        # Verify stream_id format
        assert "trace_exec1_session1" in debug_manager.manager.active_connections

    @pytest.mark.asyncio
    async def test_notify_variable_changed(self, debug_manager):
        """Test variable change notification."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await debug_manager.notify_variable_changed(
            session_id="test123",
            variable_name="counter",
            new_value=42,
            previous_value=41
        )

        assert sent_count == 1
        # Verify message structure
        call_args = mock_ws.send_text.call_args[0][0]
        import json
        message = json.loads(call_args)
        assert message["type"] == "variable_changed"
        assert message["variable"]["name"] == "counter"
        assert message["variable"]["new_value"] == 42
        assert message["variable"]["previous_value"] == 41

    @pytest.mark.asyncio
    async def test_notify_breakpoint_hit(self, debug_manager):
        """Test breakpoint hit notification."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await debug_manager.notify_breakpoint_hit(
            session_id="test123",
            breakpoint_id="bp1",
            node_id="node1",
            hit_count=5
        )

        assert sent_count == 1

    @pytest.mark.asyncio
    async def test_notify_session_paused(self, debug_manager):
        """Test session paused notification."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await debug_manager.notify_session_paused(
            session_id="test123",
            reason="user_action",
            node_id="node1"
        )

        assert sent_count == 1

    @pytest.mark.asyncio
    async def test_notify_session_resumed(self, debug_manager):
        """Test session resumed notification."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await debug_manager.notify_session_resumed("test123")

        assert sent_count == 1

    @pytest.mark.asyncio
    async def test_notify_step_completed(self, debug_manager):
        """Test step completed notification."""
        mock_ws = Mock(spec=WebSocket)
        mock_ws.accept = AsyncMock()
        mock_ws.send_text = AsyncMock()

        await debug_manager.manager.connect(mock_ws, stream_id="debug_session_test123")

        sent_count = await debug_manager.notify_step_completed(
            session_id="test123",
            action="step_over",
            step_number=5,
            node_id="node1"
        )

        assert sent_count == 1

# ============================================================================
# Singleton Tests
# ============================================================================

class TestWebSocketManagerSingleton:
    """Test WebSocket manager singleton pattern."""

    def test_get_websocket_manager_returns_singleton(self):
        """Test that get_websocket_manager returns same instance."""
        # Reset singleton
        import core.websocket_manager as ws_module
        ws_module._manager = None

        manager1 = get_websocket_manager()
        manager2 = get_websocket_manager()

        assert manager1 is manager2

    def test_get_debugging_websocket_manager_returns_singleton(self):
        """Test that get_debugging_websocket_manager returns same instance."""
        # Reset singletons
        import core.websocket_manager as ws_module
        ws_module._manager = None
        ws_module._debug_manager = None

        debug_mgr1 = get_debugging_websocket_manager()
        debug_mgr2 = get_debugging_websocket_manager()

        assert debug_mgr1 is debug_mgr2