Adding tests

tests/unit/test_base/test_publisher_graceful_shutdown.py

@@ -0,0 +1,330 @@
+"""Unit tests for Publisher graceful shutdown functionality."""
+
+import pytest
+import asyncio
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+from trustgraph.base.publisher import Publisher
+
+
+@pytest.fixture
+def mock_pulsar_client():
+    """Mock Pulsar client for testing."""
+    client = MagicMock()
+    producer = AsyncMock()
+    producer.send = MagicMock()
+    producer.flush = MagicMock()
+    producer.close = MagicMock()
+    client.create_producer.return_value = producer
+    return client
+
+
+@pytest.fixture
+def publisher(mock_pulsar_client):
+    """Create Publisher instance for testing."""
+    return Publisher(
+        client=mock_pulsar_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0
+    )
+
+
+@pytest.mark.asyncio
+async def test_publisher_queue_drain():
+    """Verify Publisher drains queue on shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic", 
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0  # Shorter timeout for testing
+    )
+    
+    # Don't start the actual run loop - just test the drain logic
+    # Fill queue with messages directly
+    for i in range(5):
+        await publisher.q.put((f"id-{i}", {"data": i}))
+    
+    # Verify queue has messages
+    assert not publisher.q.empty()
+    
+    # Mock the producer creation in run() method by patching
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a realistic run implementation that processes the queue
+        async def mock_run_impl():
+            # Simulate the actual run logic for drain
+            producer = mock_producer
+            while not publisher.q.empty():
+                try:
+                    id, item = await asyncio.wait_for(publisher.q.get(), timeout=0.1)
+                    producer.send(item, {"id": id})
+                except asyncio.TimeoutError:
+                    break
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_impl
+        
+        # Start and stop publisher
+        await publisher.start()
+        await publisher.stop()
+    
+    # Verify all messages were sent
+    assert publisher.q.empty()
+    assert mock_producer.send.call_count == 5
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_rejects_messages_during_drain():
+    """Verify Publisher rejects new messages during shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Don't start the actual run loop
+    # Add one message directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    
+    # Start shutdown process manually
+    publisher.running = False
+    publisher.draining = True
+    
+    # Try to send message during drain
+    with pytest.raises(RuntimeError, match="Publisher is shutting down"):
+        await publisher.send("id-2", {"data": 2})
+
+
+@pytest.mark.asyncio
+async def test_publisher_drain_timeout():
+    """Verify Publisher respects drain timeout."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=0.2  # Short timeout for testing
+    )
+    
+    # Fill queue with many messages directly
+    for i in range(10):
+        await publisher.q.put((f"id-{i}", {"data": i}))
+    
+    # Mock slow message processing
+    def slow_send(*args, **kwargs):
+        time.sleep(0.1)  # Simulate slow send
+    
+    mock_producer.send.side_effect = slow_send
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a run implementation that respects timeout
+        async def mock_run_with_timeout():
+            producer = mock_producer
+            end_time = time.time() + publisher.drain_timeout
+            
+            while not publisher.q.empty() and time.time() < end_time:
+                try:
+                    id, item = await asyncio.wait_for(publisher.q.get(), timeout=0.05)
+                    producer.send(item, {"id": id})
+                except asyncio.TimeoutError:
+                    break
+            
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_with_timeout
+        
+        start_time = time.time()
+        await publisher.start()
+        await publisher.stop()
+        end_time = time.time()
+    
+    # Should timeout quickly
+    assert end_time - start_time < 1.0
+    
+    # Should have called flush and close even with timeout
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_successful_drain():
+    """Verify Publisher drains successfully under normal conditions."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0
+    )
+    
+    # Add messages directly to queue
+    messages = []
+    for i in range(3):
+        msg = {"data": i}
+        await publisher.q.put((f"id-{i}", msg))
+        messages.append(msg)
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Create a successful drain implementation
+        async def mock_successful_drain():
+            producer = mock_producer
+            processed = []
+            
+            while not publisher.q.empty():
+                id, item = await publisher.q.get()
+                producer.send(item, {"id": id})
+                processed.append((id, item))
+            
+            producer.flush()
+            producer.close()
+            return processed
+        
+        mock_run.side_effect = mock_successful_drain
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # All messages should be sent
+    assert publisher.q.empty()
+    assert mock_producer.send.call_count == 3
+    
+    # Verify correct messages were sent
+    sent_calls = mock_producer.send.call_args_list
+    for i, call in enumerate(sent_calls):
+        args, kwargs = call
+        assert args[0] == {"data": i}  # message content
+        # Note: kwargs format depends on how send was called in mock
+        # Just verify message was sent with correct content
+
+
+@pytest.mark.asyncio
+async def test_publisher_state_transitions():
+    """Test Publisher state transitions during graceful shutdown."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Initial state
+    assert publisher.running is True
+    assert publisher.draining is False
+    
+    # Add message directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Mock run that simulates state transitions
+        async def mock_run_with_states():
+            # Simulate drain process
+            publisher.running = False
+            publisher.draining = True
+            
+            # Process messages
+            while not publisher.q.empty():
+                id, item = await publisher.q.get()
+                mock_producer.send(item, {"id": id})
+            
+            # Complete drain
+            publisher.draining = False
+            mock_producer.flush()
+            mock_producer.close()
+        
+        mock_run.side_effect = mock_run_with_states
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # Should have completed all state transitions
+    assert publisher.running is False
+    assert publisher.draining is False
+    mock_producer.send.assert_called_once()
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_publisher_exception_handling():
+    """Test Publisher handles exceptions during drain gracefully."""
+    mock_client = MagicMock()
+    mock_producer = MagicMock()
+    mock_client.create_producer.return_value = mock_producer
+    
+    # Mock producer.send to raise exception on second call
+    call_count = 0
+    def failing_send(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 2:
+            raise Exception("Send failed")
+    
+    mock_producer.send.side_effect = failing_send
+    
+    publisher = Publisher(
+        client=mock_client,
+        topic="test-topic",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    # Add messages directly
+    await publisher.q.put(("id-1", {"data": 1}))
+    await publisher.q.put(("id-2", {"data": 2}))
+    
+    with patch.object(publisher, 'run') as mock_run:
+        # Mock run that handles exceptions gracefully
+        async def mock_run_with_exceptions():
+            producer = mock_producer
+            
+            while not publisher.q.empty():
+                try:
+                    id, item = await publisher.q.get()
+                    producer.send(item, {"id": id})
+                except Exception as e:
+                    # Log exception but continue processing
+                    continue
+            
+            # Always call flush and close
+            producer.flush()
+            producer.close()
+        
+        mock_run.side_effect = mock_run_with_exceptions
+        
+        await publisher.start()
+        await publisher.stop()
+    
+    # Should have attempted to send both messages
+    assert mock_producer.send.call_count == 2
+    mock_producer.flush.assert_called_once()
+    mock_producer.close.assert_called_once()

tests/unit/test_base/test_subscriber_graceful_shutdown.py

@@ -0,0 +1,315 @@
+"""Unit tests for Subscriber graceful shutdown functionality."""
+
+import pytest
+import asyncio
+import uuid
+from unittest.mock import AsyncMock, MagicMock, patch
+from trustgraph.base.subscriber import Subscriber
+
+
+@pytest.fixture
+def mock_pulsar_client():
+    """Mock Pulsar client for testing."""
+    client = MagicMock()
+    consumer = MagicMock()
+    consumer.receive = MagicMock()
+    consumer.acknowledge = MagicMock()
+    consumer.negative_acknowledge = MagicMock()
+    consumer.pause_message_listener = MagicMock()
+    consumer.unsubscribe = MagicMock()
+    consumer.close = MagicMock()
+    client.subscribe.return_value = consumer
+    return client
+
+
+@pytest.fixture
+def subscriber(mock_pulsar_client):
+    """Create Subscriber instance for testing."""
+    return Subscriber(
+        client=mock_pulsar_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=2.0,
+        backpressure_strategy="block"
+    )
+
+
+def create_mock_message(message_id="test-id", data=None):
+    """Create a mock Pulsar message."""
+    msg = MagicMock()
+    msg.properties.return_value = {"id": message_id}
+    msg.value.return_value = data or {"test": "data"}
+    return msg
+
+
+@pytest.mark.asyncio
+async def test_subscriber_deferred_acknowledgment_success():
+    """Verify Subscriber only acks on successful delivery."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription", 
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        backpressure_strategy="block"
+    )
+    
+    # Create queue for subscription
+    queue = await subscriber.subscribe("test-queue")
+    
+    # Create mock message
+    msg = create_mock_message("msg-1", {"data": "test"})
+    
+    # Process message
+    await subscriber._process_message(msg)
+    
+    # Should acknowledge successful delivery
+    mock_consumer.acknowledge.assert_called_once_with(msg)
+    mock_consumer.negative_acknowledge.assert_not_called()
+    
+    # Message should be in queue
+    assert not queue.empty()
+    received_msg = await queue.get()
+    assert received_msg == {"data": "test"}
+
+
+@pytest.mark.asyncio
+async def test_subscriber_deferred_acknowledgment_failure():
+    """Verify Subscriber negative acks on delivery failure."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer", 
+        schema=dict,
+        max_size=1,  # Very small queue
+        backpressure_strategy="drop_new"
+    )
+    
+    # Create queue and fill it
+    queue = await subscriber.subscribe("test-queue")
+    await queue.put({"existing": "data"})
+    
+    # Create mock message - should be dropped
+    msg = create_mock_message("msg-1", {"data": "test"})
+    
+    # Process message (should fail due to full queue + drop_new strategy)
+    await subscriber._process_message(msg)
+    
+    # Should negative acknowledge failed delivery
+    mock_consumer.negative_acknowledge.assert_called_once_with(msg)
+    mock_consumer.acknowledge.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_backpressure_strategies():
+    """Test different backpressure strategies."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    # Test drop_oldest strategy
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic", 
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=2,
+        backpressure_strategy="drop_oldest"
+    )
+    
+    queue = await subscriber.subscribe("test-queue")
+    
+    # Fill queue
+    await queue.put({"data": "old1"})
+    await queue.put({"data": "old2"})
+    
+    # Add new message (should drop oldest)
+    msg = create_mock_message("msg-1", {"data": "new"})
+    await subscriber._process_message(msg)
+    
+    # Should acknowledge delivery
+    mock_consumer.acknowledge.assert_called_once_with(msg)
+    
+    # Queue should have new message (old one dropped)
+    messages = []
+    while not queue.empty():
+        messages.append(await queue.get())
+    
+    # Should contain old2 and new (old1 was dropped)
+    assert len(messages) == 2
+    assert {"data": "new"} in messages
+
+
+@pytest.mark.asyncio
+async def test_subscriber_graceful_shutdown():
+    """Test Subscriber graceful shutdown with queue draining."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=1.0
+    )
+    
+    await subscriber.start()
+    
+    # Create subscription with messages
+    queue = await subscriber.subscribe("test-queue")
+    await queue.put({"data": "msg1"})
+    await queue.put({"data": "msg2"})
+    
+    # Initial state
+    assert subscriber.running is True
+    assert subscriber.draining is False
+    
+    # Start shutdown
+    stop_task = asyncio.create_task(subscriber.stop())
+    
+    # Allow brief processing
+    await asyncio.sleep(0.1)
+    
+    # Should be in drain state
+    assert subscriber.running is False
+    assert subscriber.draining is True
+    
+    # Should pause message listener
+    mock_consumer.pause_message_listener.assert_called_once()
+    
+    # Complete shutdown
+    await stop_task
+    
+    # Should have cleaned up
+    mock_consumer.unsubscribe.assert_called_once()
+    mock_consumer.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_drain_timeout():
+    """Test Subscriber respects drain timeout."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription", 
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10,
+        drain_timeout=0.1  # Very short timeout
+    )
+    
+    await subscriber.start()
+    
+    # Create subscription with many messages
+    queue = await subscriber.subscribe("test-queue")
+    for i in range(20):
+        await queue.put({"data": f"msg{i}"})
+    
+    import time
+    start_time = time.time()
+    await subscriber.stop()
+    end_time = time.time()
+    
+    # Should timeout quickly
+    assert end_time - start_time < 1.0
+    
+    # Queue should still have messages (drain timed out)
+    assert not queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_subscriber_pending_acks_cleanup():
+    """Test Subscriber cleans up pending acknowledgments on shutdown."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10
+    )
+    
+    await subscriber.start()
+    
+    # Add pending acknowledgments manually (simulating in-flight messages)
+    msg1 = create_mock_message("msg-1")
+    msg2 = create_mock_message("msg-2") 
+    subscriber.pending_acks["ack-1"] = msg1
+    subscriber.pending_acks["ack-2"] = msg2
+    
+    # Stop subscriber
+    await subscriber.stop()
+    
+    # Should negative acknowledge pending messages
+    assert mock_consumer.negative_acknowledge.call_count == 2
+    mock_consumer.negative_acknowledge.assert_any_call(msg1)
+    mock_consumer.negative_acknowledge.assert_any_call(msg2)
+    
+    # Pending acks should be cleared
+    assert len(subscriber.pending_acks) == 0
+
+
+@pytest.mark.asyncio
+async def test_subscriber_multiple_subscribers():
+    """Test Subscriber with multiple concurrent subscribers."""
+    mock_client = MagicMock()
+    mock_consumer = MagicMock()
+    mock_client.subscribe.return_value = mock_consumer
+    
+    subscriber = Subscriber(
+        client=mock_client,
+        topic="test-topic",
+        subscription="test-subscription",
+        consumer_name="test-consumer",
+        schema=dict,
+        max_size=10
+    )
+    
+    # Create multiple subscriptions
+    queue1 = await subscriber.subscribe("queue-1")
+    queue2 = await subscriber.subscribe("queue-2")
+    queue_all = await subscriber.subscribe_all("queue-all")
+    
+    # Process message
+    msg = create_mock_message("msg-1", {"data": "broadcast"})
+    await subscriber._process_message(msg)
+    
+    # Should acknowledge (successful delivery to all queues)
+    mock_consumer.acknowledge.assert_called_once_with(msg)
+    
+    # Message should be in specific queue (queue-1) and broadcast queue
+    assert not queue1.empty()
+    assert queue2.empty()  # No message for queue-2
+    assert not queue_all.empty()
+    
+    # Verify message content
+    msg1 = await queue1.get()
+    msg_all = await queue_all.get()
+    assert msg1 == {"data": "broadcast"}
+    assert msg_all == {"data": "broadcast"}

tests/unit/test_gateway/test_socket_graceful_shutdown.py

@@ -0,0 +1,324 @@
+"""Unit tests for SocketEndpoint graceful shutdown functionality."""
+
+import pytest
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+from aiohttp import web, WSMsgType
+from trustgraph.gateway.endpoint.socket import SocketEndpoint
+from trustgraph.gateway.running import Running
+
+
+@pytest.fixture
+def mock_auth():
+    """Mock authentication service."""
+    auth = MagicMock()
+    auth.permitted.return_value = True
+    return auth
+
+
+@pytest.fixture
+def mock_dispatcher_factory():
+    """Mock dispatcher factory function."""
+    async def dispatcher_factory(ws, running, match_info):
+        dispatcher = AsyncMock()
+        dispatcher.run = AsyncMock()
+        dispatcher.receive = AsyncMock()
+        dispatcher.destroy = AsyncMock()
+        return dispatcher
+    
+    return dispatcher_factory
+
+
+@pytest.fixture
+def socket_endpoint(mock_auth, mock_dispatcher_factory):
+    """Create SocketEndpoint for testing."""
+    return SocketEndpoint(
+        endpoint_path="/test-socket",
+        auth=mock_auth,
+        dispatcher=mock_dispatcher_factory
+    )
+
+
+@pytest.fixture
+def mock_websocket():
+    """Mock websocket response."""
+    ws = AsyncMock(spec=web.WebSocketResponse)
+    ws.prepare = AsyncMock()
+    ws.close = AsyncMock()
+    ws.closed = False
+    return ws
+
+
+@pytest.fixture
+def mock_request():
+    """Mock HTTP request."""
+    request = MagicMock()
+    request.query = {"token": "test-token"}
+    request.match_info = {}
+    return request
+
+
+@pytest.mark.asyncio
+async def test_listener_graceful_shutdown_on_close():
+    """Test listener handles websocket close gracefully."""
+    socket_endpoint = SocketEndpoint("/test", MagicMock(), AsyncMock())
+    
+    # Mock websocket that closes after one message
+    ws = AsyncMock()
+    
+    # Create async iterator that yields one message then closes
+    async def mock_iterator():
+        # Yield normal message
+        msg = MagicMock()
+        msg.type = WSMsgType.TEXT
+        yield msg
+        
+        # Yield close message 
+        close_msg = MagicMock()
+        close_msg.type = WSMsgType.CLOSE
+        yield close_msg
+    
+    ws.__aiter__ = mock_iterator
+    
+    dispatcher = AsyncMock()
+    running = Running()
+    
+    with patch('asyncio.sleep') as mock_sleep:
+        await socket_endpoint.listener(ws, dispatcher, running)
+    
+    # Should have processed one message
+    dispatcher.receive.assert_called_once()
+    
+    # Should have initiated graceful shutdown
+    assert running.get() is False
+    
+    # Should have slept for grace period
+    mock_sleep.assert_called_once_with(1.0)
+
+
+@pytest.mark.asyncio
+async def test_handle_normal_flow():
+    """Test normal websocket handling flow."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = True
+    
+    dispatcher_created = False
+    async def mock_dispatcher_factory(ws, running, match_info):
+        nonlocal dispatcher_created
+        dispatcher_created = True
+        dispatcher = AsyncMock()
+        dispatcher.destroy = AsyncMock()
+        return dispatcher
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, mock_dispatcher_factory)
+    
+    request = MagicMock()
+    request.query = {"token": "valid-token"}
+    request.match_info = {}
+    
+    with patch('aiohttp.web.WebSocketResponse') as mock_ws_class:
+        mock_ws = AsyncMock()
+        mock_ws.prepare = AsyncMock()
+        mock_ws.close = AsyncMock()
+        mock_ws.closed = False
+        mock_ws_class.return_value = mock_ws
+        
+        with patch('asyncio.TaskGroup') as mock_task_group:
+            # Mock task group context manager
+            mock_tg = AsyncMock()
+            mock_tg.__aenter__ = AsyncMock(return_value=mock_tg)
+            mock_tg.__aexit__ = AsyncMock(return_value=None)
+            mock_tg.create_task = MagicMock(return_value=AsyncMock())
+            mock_task_group.return_value = mock_tg
+            
+            result = await socket_endpoint.handle(request)
+    
+    # Should have created dispatcher
+    assert dispatcher_created is True
+    
+    # Should return websocket
+    assert result == mock_ws
+
+
+@pytest.mark.asyncio
+async def test_handle_exception_group_cleanup():
+    """Test exception group triggers dispatcher cleanup."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = True
+    
+    mock_dispatcher = AsyncMock()
+    mock_dispatcher.destroy = AsyncMock()
+    
+    async def mock_dispatcher_factory(ws, running, match_info):
+        return mock_dispatcher
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, mock_dispatcher_factory)
+    
+    request = MagicMock()
+    request.query = {"token": "valid-token"}
+    request.match_info = {}
+    
+    # Mock TaskGroup to raise ExceptionGroup
+    class TestException(Exception):
+        pass
+    
+    exception_group = ExceptionGroup("Test exceptions", [TestException("test")])
+    
+    with patch('aiohttp.web.WebSocketResponse') as mock_ws_class:
+        mock_ws = AsyncMock()
+        mock_ws.prepare = AsyncMock()
+        mock_ws.close = AsyncMock() 
+        mock_ws.closed = False
+        mock_ws_class.return_value = mock_ws
+        
+        with patch('asyncio.TaskGroup') as mock_task_group:
+            mock_tg = AsyncMock()
+            mock_tg.__aenter__ = AsyncMock(side_effect=exception_group)
+            mock_tg.__aexit__ = AsyncMock(return_value=None)
+            mock_task_group.return_value = mock_tg
+            
+            with patch('asyncio.wait_for') as mock_wait_for:
+                mock_wait_for.return_value = None
+                
+                result = await socket_endpoint.handle(request)
+    
+    # Should have attempted graceful cleanup
+    mock_wait_for.assert_called_once()
+    
+    # Should have called destroy in finally block
+    assert mock_dispatcher.destroy.call_count >= 1
+    
+    # Should have closed websocket
+    mock_ws.close.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_handle_dispatcher_cleanup_timeout():
+    """Test dispatcher cleanup with timeout."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = True
+    
+    # Mock dispatcher that takes long to destroy
+    mock_dispatcher = AsyncMock()
+    mock_dispatcher.destroy = AsyncMock()
+    
+    async def mock_dispatcher_factory(ws, running, match_info):
+        return mock_dispatcher
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, mock_dispatcher_factory)
+    
+    request = MagicMock()
+    request.query = {"token": "valid-token"}
+    request.match_info = {}
+    
+    # Mock TaskGroup to raise exception
+    exception_group = ExceptionGroup("Test", [Exception("test")])
+    
+    with patch('aiohttp.web.WebSocketResponse') as mock_ws_class:
+        mock_ws = AsyncMock()
+        mock_ws.prepare = AsyncMock()
+        mock_ws.close = AsyncMock()
+        mock_ws.closed = False
+        mock_ws_class.return_value = mock_ws
+        
+        with patch('asyncio.TaskGroup') as mock_task_group:
+            mock_tg = AsyncMock()
+            mock_tg.__aenter__ = AsyncMock(side_effect=exception_group)
+            mock_tg.__aexit__ = AsyncMock(return_value=None)
+            mock_task_group.return_value = mock_tg
+            
+            # Mock asyncio.wait_for to raise TimeoutError
+            with patch('asyncio.wait_for') as mock_wait_for:
+                mock_wait_for.side_effect = asyncio.TimeoutError("Cleanup timeout")
+                
+                result = await socket_endpoint.handle(request)
+    
+    # Should have attempted cleanup with timeout
+    mock_wait_for.assert_called_once_with(
+        mock_dispatcher.destroy(),
+        timeout=5.0
+    )
+    
+    # Should still call destroy in finally block
+    assert mock_dispatcher.destroy.call_count >= 1
+
+
+@pytest.mark.asyncio
+async def test_handle_unauthorized_request():
+    """Test handling of unauthorized requests."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = False  # Unauthorized
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, AsyncMock())
+    
+    request = MagicMock()
+    request.query = {"token": "invalid-token"}
+    
+    result = await socket_endpoint.handle(request)
+    
+    # Should return HTTP 401
+    assert isinstance(result, web.HTTPUnauthorized)
+    
+    # Should have checked permission
+    mock_auth.permitted.assert_called_once_with("invalid-token", "socket")
+
+
+@pytest.mark.asyncio
+async def test_handle_missing_token():
+    """Test handling of requests with missing token."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = False
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, AsyncMock())
+    
+    request = MagicMock()
+    request.query = {}  # No token
+    
+    result = await socket_endpoint.handle(request)
+    
+    # Should return HTTP 401
+    assert isinstance(result, web.HTTPUnauthorized)
+    
+    # Should have checked permission with empty token
+    mock_auth.permitted.assert_called_once_with("", "socket")
+
+
+@pytest.mark.asyncio
+async def test_handle_websocket_already_closed():
+    """Test handling when websocket is already closed."""
+    mock_auth = MagicMock()
+    mock_auth.permitted.return_value = True
+    
+    mock_dispatcher = AsyncMock()
+    mock_dispatcher.destroy = AsyncMock()
+    
+    async def mock_dispatcher_factory(ws, running, match_info):
+        return mock_dispatcher
+    
+    socket_endpoint = SocketEndpoint("/test", mock_auth, mock_dispatcher_factory)
+    
+    request = MagicMock()
+    request.query = {"token": "valid-token"}
+    request.match_info = {}
+    
+    with patch('aiohttp.web.WebSocketResponse') as mock_ws_class:
+        mock_ws = AsyncMock()
+        mock_ws.prepare = AsyncMock()
+        mock_ws.close = AsyncMock()
+        mock_ws.closed = True  # Already closed
+        mock_ws_class.return_value = mock_ws
+        
+        with patch('asyncio.TaskGroup') as mock_task_group:
+            mock_tg = AsyncMock()
+            mock_tg.__aenter__ = AsyncMock(return_value=mock_tg)
+            mock_tg.__aexit__ = AsyncMock(return_value=None)
+            mock_tg.create_task = MagicMock(return_value=AsyncMock())
+            mock_task_group.return_value = mock_tg
+            
+            result = await socket_endpoint.handle(request)
+    
+    # Should still have called destroy
+    mock_dispatcher.destroy.assert_called()
+    
+    # Should not attempt to close already closed websocket
+    mock_ws.close.assert_not_called()  # Not called in finally since ws.closed = True