mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
96c2b73457
* 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
136 lines
No EOL
4.6 KiB
Python
136 lines
No EOL
4.6 KiB
Python
"""
|
|
Tests for Gateway Socket Endpoint
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, AsyncMock
|
|
from aiohttp import WSMsgType
|
|
|
|
from trustgraph.gateway.endpoint.socket import SocketEndpoint
|
|
|
|
|
|
class TestSocketEndpoint:
|
|
"""Test cases for SocketEndpoint class"""
|
|
|
|
def test_socket_endpoint_initialization(self):
|
|
"""Test SocketEndpoint initialization"""
|
|
mock_auth = MagicMock()
|
|
mock_dispatcher = MagicMock()
|
|
|
|
endpoint = SocketEndpoint(
|
|
endpoint_path="/api/socket",
|
|
auth=mock_auth,
|
|
dispatcher=mock_dispatcher
|
|
)
|
|
|
|
assert endpoint.path == "/api/socket"
|
|
assert endpoint.auth == mock_auth
|
|
assert endpoint.dispatcher == mock_dispatcher
|
|
assert endpoint.operation == "socket"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_method(self):
|
|
"""Test SocketEndpoint worker method"""
|
|
mock_auth = MagicMock()
|
|
mock_dispatcher = AsyncMock()
|
|
|
|
endpoint = SocketEndpoint("/api/socket", mock_auth, mock_dispatcher)
|
|
|
|
mock_ws = MagicMock()
|
|
mock_running = MagicMock()
|
|
|
|
# Call worker method
|
|
await endpoint.worker(mock_ws, mock_dispatcher, mock_running)
|
|
|
|
# Verify dispatcher.run was called
|
|
mock_dispatcher.run.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listener_method_with_text_message(self):
|
|
"""Test SocketEndpoint listener method with text message"""
|
|
mock_auth = MagicMock()
|
|
mock_dispatcher = AsyncMock()
|
|
|
|
endpoint = SocketEndpoint("/api/socket", mock_auth, mock_dispatcher)
|
|
|
|
# Mock websocket with text message
|
|
mock_msg = MagicMock()
|
|
mock_msg.type = WSMsgType.TEXT
|
|
|
|
# Create async iterator for websocket
|
|
async def async_iter():
|
|
yield mock_msg
|
|
|
|
mock_ws = AsyncMock()
|
|
mock_ws.__aiter__ = lambda self: async_iter()
|
|
mock_ws.closed = False # Set closed attribute
|
|
mock_running = MagicMock()
|
|
|
|
# Call listener method
|
|
await endpoint.listener(mock_ws, mock_dispatcher, mock_running)
|
|
|
|
# Verify dispatcher.receive was called with the message
|
|
mock_dispatcher.receive.assert_called_once_with(mock_msg)
|
|
# Verify cleanup methods were called
|
|
mock_running.stop.assert_called_once()
|
|
mock_ws.close.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listener_method_with_binary_message(self):
|
|
"""Test SocketEndpoint listener method with binary message"""
|
|
mock_auth = MagicMock()
|
|
mock_dispatcher = AsyncMock()
|
|
|
|
endpoint = SocketEndpoint("/api/socket", mock_auth, mock_dispatcher)
|
|
|
|
# Mock websocket with binary message
|
|
mock_msg = MagicMock()
|
|
mock_msg.type = WSMsgType.BINARY
|
|
|
|
# Create async iterator for websocket
|
|
async def async_iter():
|
|
yield mock_msg
|
|
|
|
mock_ws = AsyncMock()
|
|
mock_ws.__aiter__ = lambda self: async_iter()
|
|
mock_ws.closed = False # Set closed attribute
|
|
mock_running = MagicMock()
|
|
|
|
# Call listener method
|
|
await endpoint.listener(mock_ws, mock_dispatcher, mock_running)
|
|
|
|
# Verify dispatcher.receive was called with the message
|
|
mock_dispatcher.receive.assert_called_once_with(mock_msg)
|
|
# Verify cleanup methods were called
|
|
mock_running.stop.assert_called_once()
|
|
mock_ws.close.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_listener_method_with_close_message(self):
|
|
"""Test SocketEndpoint listener method with close message"""
|
|
mock_auth = MagicMock()
|
|
mock_dispatcher = AsyncMock()
|
|
|
|
endpoint = SocketEndpoint("/api/socket", mock_auth, mock_dispatcher)
|
|
|
|
# Mock websocket with close message
|
|
mock_msg = MagicMock()
|
|
mock_msg.type = WSMsgType.CLOSE
|
|
|
|
# Create async iterator for websocket
|
|
async def async_iter():
|
|
yield mock_msg
|
|
|
|
mock_ws = AsyncMock()
|
|
mock_ws.__aiter__ = lambda self: async_iter()
|
|
mock_ws.closed = False # Set closed attribute
|
|
mock_running = MagicMock()
|
|
|
|
# Call listener method
|
|
await endpoint.listener(mock_ws, mock_dispatcher, mock_running)
|
|
|
|
# Verify dispatcher.receive was NOT called for close message
|
|
mock_dispatcher.receive.assert_not_called()
|
|
# Verify cleanup methods were called after break
|
|
mock_running.stop.assert_called_once()
|
|
mock_ws.close.assert_called_once() |