mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
365 lines
12 KiB
Python
365 lines
12 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import hashlib
|
|
from io import StringIO
|
|
import json
|
|
import sys
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import Mock
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
# Skip all tests in this module if Python version is less than 3.10
|
|
pytestmark = pytest.mark.skipif(
|
|
sys.version_info < (3, 10), reason="MCP tool requires Python 3.10+"
|
|
)
|
|
|
|
# Import dependencies with version checking
|
|
try:
|
|
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
|
|
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource
|
|
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
|
|
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
|
|
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
|
|
except ImportError as e:
|
|
if sys.version_info < (3, 10):
|
|
# Create dummy classes to prevent NameError during test collection
|
|
# Tests will be skipped anyway due to pytestmark
|
|
class DummyClass:
|
|
pass
|
|
|
|
MCPSessionManager = DummyClass
|
|
retry_on_closed_resource = lambda x: x
|
|
SseConnectionParams = DummyClass
|
|
StdioConnectionParams = DummyClass
|
|
StreamableHTTPConnectionParams = DummyClass
|
|
else:
|
|
raise e
|
|
|
|
# Import real MCP classes
|
|
try:
|
|
from mcp import StdioServerParameters
|
|
except ImportError:
|
|
# Create a mock if MCP is not available
|
|
class StdioServerParameters:
|
|
|
|
def __init__(self, command="test_command", args=None):
|
|
self.command = command
|
|
self.args = args or []
|
|
|
|
|
|
class MockClientSession:
|
|
"""Mock ClientSession for testing."""
|
|
|
|
def __init__(self):
|
|
self._read_stream = Mock()
|
|
self._write_stream = Mock()
|
|
self._read_stream._closed = False
|
|
self._write_stream._closed = False
|
|
self.initialize = AsyncMock()
|
|
|
|
|
|
class MockAsyncExitStack:
|
|
"""Mock AsyncExitStack for testing."""
|
|
|
|
def __init__(self):
|
|
self.aclose = AsyncMock()
|
|
self.enter_async_context = AsyncMock()
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
class TestMCPSessionManager:
|
|
"""Test suite for MCPSessionManager class."""
|
|
|
|
def setup_method(self):
|
|
"""Set up test fixtures."""
|
|
self.mock_stdio_params = StdioServerParameters(
|
|
command="test_command", args=[]
|
|
)
|
|
self.mock_stdio_connection_params = StdioConnectionParams(
|
|
server_params=self.mock_stdio_params, timeout=5.0
|
|
)
|
|
|
|
def test_init_with_stdio_server_parameters(self):
|
|
"""Test initialization with StdioServerParameters (deprecated)."""
|
|
with patch(
|
|
"google.adk.tools.mcp_tool.mcp_session_manager.logger"
|
|
) as mock_logger:
|
|
manager = MCPSessionManager(self.mock_stdio_params)
|
|
|
|
# Should log deprecation warning
|
|
mock_logger.warning.assert_called_once()
|
|
assert "StdioServerParameters is not recommended" in str(
|
|
mock_logger.warning.call_args
|
|
)
|
|
|
|
# Should convert to StdioConnectionParams
|
|
assert isinstance(manager._connection_params, StdioConnectionParams)
|
|
assert manager._connection_params.server_params == self.mock_stdio_params
|
|
assert manager._connection_params.timeout == 5
|
|
|
|
def test_init_with_stdio_connection_params(self):
|
|
"""Test initialization with StdioConnectionParams."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
assert manager._connection_params == self.mock_stdio_connection_params
|
|
assert manager._errlog == sys.stderr
|
|
assert manager._sessions == {}
|
|
|
|
def test_init_with_sse_connection_params(self):
|
|
"""Test initialization with SseConnectionParams."""
|
|
sse_params = SseConnectionParams(
|
|
url="https://example.com/mcp",
|
|
headers={"Authorization": "Bearer token"},
|
|
timeout=10.0,
|
|
)
|
|
manager = MCPSessionManager(sse_params)
|
|
|
|
assert manager._connection_params == sse_params
|
|
|
|
def test_init_with_streamable_http_params(self):
|
|
"""Test initialization with StreamableHTTPConnectionParams."""
|
|
http_params = StreamableHTTPConnectionParams(
|
|
url="https://example.com/mcp", timeout=15.0
|
|
)
|
|
manager = MCPSessionManager(http_params)
|
|
|
|
assert manager._connection_params == http_params
|
|
|
|
def test_generate_session_key_stdio(self):
|
|
"""Test session key generation for stdio connections."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# For stdio, headers should be ignored and return constant key
|
|
key1 = manager._generate_session_key({"Authorization": "Bearer token"})
|
|
key2 = manager._generate_session_key(None)
|
|
|
|
assert key1 == "stdio_session"
|
|
assert key2 == "stdio_session"
|
|
assert key1 == key2
|
|
|
|
def test_generate_session_key_sse(self):
|
|
"""Test session key generation for SSE connections."""
|
|
sse_params = SseConnectionParams(url="https://example.com/mcp")
|
|
manager = MCPSessionManager(sse_params)
|
|
|
|
headers1 = {"Authorization": "Bearer token1"}
|
|
headers2 = {"Authorization": "Bearer token2"}
|
|
|
|
key1 = manager._generate_session_key(headers1)
|
|
key2 = manager._generate_session_key(headers2)
|
|
key3 = manager._generate_session_key(headers1)
|
|
|
|
# Different headers should generate different keys
|
|
assert key1 != key2
|
|
# Same headers should generate same key
|
|
assert key1 == key3
|
|
|
|
# Should be deterministic hash
|
|
headers_json = json.dumps(headers1, sort_keys=True)
|
|
expected_hash = hashlib.md5(headers_json.encode()).hexdigest()
|
|
assert key1 == f"session_{expected_hash}"
|
|
|
|
def test_merge_headers_stdio(self):
|
|
"""Test header merging for stdio connections."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# Stdio connections don't support headers
|
|
headers = manager._merge_headers({"Authorization": "Bearer token"})
|
|
assert headers is None
|
|
|
|
def test_merge_headers_sse(self):
|
|
"""Test header merging for SSE connections."""
|
|
base_headers = {"Content-Type": "application/json"}
|
|
sse_params = SseConnectionParams(
|
|
url="https://example.com/mcp", headers=base_headers
|
|
)
|
|
manager = MCPSessionManager(sse_params)
|
|
|
|
# With additional headers
|
|
additional = {"Authorization": "Bearer token"}
|
|
merged = manager._merge_headers(additional)
|
|
|
|
expected = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": "Bearer token",
|
|
}
|
|
assert merged == expected
|
|
|
|
def test_is_session_disconnected(self):
|
|
"""Test session disconnection detection."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# Create mock session
|
|
session = MockClientSession()
|
|
|
|
# Not disconnected
|
|
assert not manager._is_session_disconnected(session)
|
|
|
|
# Disconnected - read stream closed
|
|
session._read_stream._closed = True
|
|
assert manager._is_session_disconnected(session)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_stdio_new(self):
|
|
"""Test creating a new stdio session."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
mock_session = MockClientSession()
|
|
mock_exit_stack = MockAsyncExitStack()
|
|
|
|
with patch(
|
|
"google.adk.tools.mcp_tool.mcp_session_manager.stdio_client"
|
|
) as mock_stdio:
|
|
with patch(
|
|
"google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack"
|
|
) as mock_exit_stack_class:
|
|
with patch(
|
|
"google.adk.tools.mcp_tool.mcp_session_manager.ClientSession"
|
|
) as mock_session_class:
|
|
|
|
# Setup mocks
|
|
mock_exit_stack_class.return_value = mock_exit_stack
|
|
mock_stdio.return_value = AsyncMock()
|
|
mock_exit_stack.enter_async_context.side_effect = [
|
|
("read", "write"), # First call returns transports
|
|
mock_session, # Second call returns session
|
|
]
|
|
mock_session_class.return_value = mock_session
|
|
|
|
# Create session
|
|
session = await manager.create_session()
|
|
|
|
# Verify session creation
|
|
assert session == mock_session
|
|
assert len(manager._sessions) == 1
|
|
assert "stdio_session" in manager._sessions
|
|
|
|
# Verify session was initialized
|
|
mock_session.initialize.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_session_reuse_existing(self):
|
|
"""Test reusing an existing connected session."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# Create mock existing session
|
|
existing_session = MockClientSession()
|
|
existing_exit_stack = MockAsyncExitStack()
|
|
manager._sessions["stdio_session"] = (existing_session, existing_exit_stack)
|
|
|
|
# Session is connected
|
|
existing_session._read_stream._closed = False
|
|
existing_session._write_stream._closed = False
|
|
|
|
session = await manager.create_session()
|
|
|
|
# Should reuse existing session
|
|
assert session == existing_session
|
|
assert len(manager._sessions) == 1
|
|
|
|
# Should not create new session
|
|
existing_session.initialize.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_success(self):
|
|
"""Test successful cleanup of all sessions."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# Add mock sessions
|
|
session1 = MockClientSession()
|
|
exit_stack1 = MockAsyncExitStack()
|
|
session2 = MockClientSession()
|
|
exit_stack2 = MockAsyncExitStack()
|
|
|
|
manager._sessions["session1"] = (session1, exit_stack1)
|
|
manager._sessions["session2"] = (session2, exit_stack2)
|
|
|
|
await manager.close()
|
|
|
|
# All sessions should be closed
|
|
exit_stack1.aclose.assert_called_once()
|
|
exit_stack2.aclose.assert_called_once()
|
|
assert len(manager._sessions) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_with_errors(self):
|
|
"""Test cleanup when some sessions fail to close."""
|
|
manager = MCPSessionManager(self.mock_stdio_connection_params)
|
|
|
|
# Add mock sessions
|
|
session1 = MockClientSession()
|
|
exit_stack1 = MockAsyncExitStack()
|
|
exit_stack1.aclose.side_effect = Exception("Close error 1")
|
|
|
|
session2 = MockClientSession()
|
|
exit_stack2 = MockAsyncExitStack()
|
|
|
|
manager._sessions["session1"] = (session1, exit_stack1)
|
|
manager._sessions["session2"] = (session2, exit_stack2)
|
|
|
|
custom_errlog = StringIO()
|
|
manager._errlog = custom_errlog
|
|
|
|
# Should not raise exception
|
|
await manager.close()
|
|
|
|
# Good session should still be closed
|
|
exit_stack2.aclose.assert_called_once()
|
|
assert len(manager._sessions) == 0
|
|
|
|
# Error should be logged
|
|
error_output = custom_errlog.getvalue()
|
|
assert "Warning: Error during MCP session cleanup" in error_output
|
|
assert "Close error 1" in error_output
|
|
|
|
|
|
def test_retry_on_closed_resource_decorator():
|
|
"""Test the retry_on_closed_resource decorator."""
|
|
|
|
call_count = 0
|
|
|
|
@retry_on_closed_resource
|
|
async def mock_function(self):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 1:
|
|
import anyio
|
|
|
|
raise anyio.ClosedResourceError("Resource closed")
|
|
return "success"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry():
|
|
nonlocal call_count
|
|
call_count = 0
|
|
|
|
mock_self = Mock()
|
|
result = await mock_function(mock_self)
|
|
|
|
assert result == "success"
|
|
assert call_count == 2 # First call fails, second succeeds
|
|
|
|
# Run the test
|
|
import asyncio
|
|
|
|
asyncio.run(test_retry())
|