Fix import export graceful shutdown (#476)

* Tech spec for graceful shutdown

* Graceful shutdown of importers/exporters

* Update socket to include graceful shutdown orchestration

* Adding tests for conditions tracked in this PR

tests/unit/test_gateway/test_socket_graceful_shutdown.py

@@ -0,0 +1,326 @@
+"""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(self):
+        # 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
+    
+    # Set the async iterator method
+    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(return_value=mock_tg)
+            mock_tg.__aexit__ = AsyncMock(side_effect=exception_group)
+            mock_tg.create_task = MagicMock(side_effect=TestException("test"))
+            mock_task_group.return_value = mock_tg
+            
+            with patch('trustgraph.gateway.endpoint.socket.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(return_value=mock_tg)
+            mock_tg.__aexit__ = AsyncMock(side_effect=exception_group)
+            mock_tg.create_task = MagicMock(side_effect=Exception("test"))
+            mock_task_group.return_value = mock_tg
+            
+            # Mock asyncio.wait_for to raise TimeoutError
+            with patch('trustgraph.gateway.endpoint.socket.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()
+    # Check that timeout was passed correctly
+    assert mock_wait_for.call_args[1]['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