This commit is contained in:
Alexander Myasoedov
2025-04-03 20:43:53 +03:00
parent 47ca656d59
commit 0dc738a11e
3 changed files with 38 additions and 30 deletions
+26 -18
View File
@@ -1,10 +1,14 @@
import pytest
import asyncio
from unittest.mock import AsyncMock, patch
import pytest
from agentic_security.mcp import ClientSession
from ..agentic_security.mcp.client import run
# Fixtures
# Fixtures
@pytest.fixture(scope="session")
def event_loop():
"""Create an instance of the default event loop for each test case."""
@@ -12,46 +16,49 @@ def event_loop():
yield loop
loop.close()
@pytest.fixture
async def mock_session():
with patch('mcp.client.stdio.stdio_client') as mock_client:
with patch("mcp.client.stdio.stdio_client") as mock_client:
mock_read = AsyncMock()
mock_write = AsyncMock()
# Configures mock client such that mock responses are returned
# Configures mock client such that mock responses are returned
mock_client.return_value.__aenter__.return_value = (mock_read, mock_write)
# Creates a mock session
mock_session = AsyncMock(spec=ClientSession)
# Expected responses
mock_session.initialize = AsyncMock()
mock_session.list_prompts = AsyncMock(return_value=["test_prompt"])
mock_session.list_resources = AsyncMock(return_value=["test_resource"])
mock_session.list_tools = AsyncMock(return_value=["echo_tool"])
mock_session.call_tool = AsyncMock(return_value="Hello from client!")
with patch('mcp.ClientSession', return_value=mock_session):
with patch("mcp.ClientSession", return_value=mock_session):
yield mock_session
# Tests
# Tests
@pytest.mark.asyncio
async def test_initialization(mock_session):
"""Test initialization success and failure cases"""
# Test initialization success case
# Test initialization success case
await run()
mock_session.initialize.assert_called_once()
# Resetting the mock to test for failure case
# Resetting the mock to test for failure case
mock_session.initialize.reset_mock()
mock_session.initialize.side_effect = ConnectionError("Failed to connect")
# Test connection error
with pytest.raises(ConnectionError):
await run()
@pytest.mark.asyncio
async def test_list_resources(mock_session):
"""Test listing available resources"""
@@ -59,6 +66,7 @@ async def test_list_resources(mock_session):
mock_session.list_resources.assert_called_once()
assert await mock_session.list_resources() == ["test_resource"]
@pytest.mark.asyncio
async def test_list_tools(mock_session):
"""Test listing available tools"""
@@ -66,11 +74,11 @@ async def test_list_tools(mock_session):
mock_session.list_tools.assert_called_once()
assert await mock_session.list_tools() == ["echo_tool"]
@pytest.mark.asyncio
async def test_echo_tool(mock_session):
"""Test the echo tool functionality"""
await run()
mock_session.call_tool.assert_called_once_with(
"echo_tool",
arguments={"message": "Hello from client!"}
)
"echo_tool", arguments={"message": "Hello from client!"}
)