|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +from contextlib import asynccontextmanager |
| 4 | +from unittest.mock import AsyncMock, MagicMock, patch |
| 5 | + |
3 | 6 | import pytest |
4 | 7 |
|
5 | | -from stackone_ai.toolset import StackOneToolSet, _McpToolDefinition |
| 8 | +from stackone_ai.toolset import StackOneToolSet, _fetch_mcp_tools, _McpToolDefinition |
6 | 9 |
|
7 | 10 |
|
8 | 11 | @pytest.fixture |
@@ -318,3 +321,96 @@ def fake_fetch(_: str, headers: dict[str, str]) -> list[_McpToolDefinition]: |
318 | 321 | toolset = StackOneToolSet(api_key="test_key") |
319 | 322 | with pytest.raises(ToolsetConfigError, match="Original config error"): |
320 | 323 | toolset.fetch_tools() |
| 324 | + |
| 325 | + |
| 326 | +class TestFetchMcpToolsInternal: |
| 327 | + """Test _fetch_mcp_tools internal implementation.""" |
| 328 | + |
| 329 | + def test_fetch_mcp_tools_single_page(self): |
| 330 | + """Test fetching tools with single page response.""" |
| 331 | + # Create mock tool response |
| 332 | + mock_tool = MagicMock() |
| 333 | + mock_tool.name = "test_tool" |
| 334 | + mock_tool.description = "Test description" |
| 335 | + mock_tool.inputSchema = {"type": "object", "properties": {"id": {"type": "string"}}} |
| 336 | + |
| 337 | + mock_result = MagicMock() |
| 338 | + mock_result.tools = [mock_tool] |
| 339 | + mock_result.nextCursor = None |
| 340 | + |
| 341 | + # Create mock session |
| 342 | + mock_session = AsyncMock() |
| 343 | + mock_session.initialize = AsyncMock() |
| 344 | + mock_session.list_tools = AsyncMock(return_value=mock_result) |
| 345 | + mock_session.__aenter__ = AsyncMock(return_value=mock_session) |
| 346 | + mock_session.__aexit__ = AsyncMock(return_value=None) |
| 347 | + |
| 348 | + # Create mock streamable client |
| 349 | + @asynccontextmanager |
| 350 | + async def mock_streamable_client(endpoint, headers): |
| 351 | + yield (MagicMock(), MagicMock(), MagicMock()) |
| 352 | + |
| 353 | + # Patch at the module where imports happen |
| 354 | + with ( |
| 355 | + patch( |
| 356 | + "mcp.client.streamable_http.streamablehttp_client", |
| 357 | + side_effect=mock_streamable_client, |
| 358 | + ), |
| 359 | + patch("mcp.client.session.ClientSession", return_value=mock_session), |
| 360 | + patch("mcp.types.Implementation", MagicMock()), |
| 361 | + ): |
| 362 | + result = _fetch_mcp_tools("https://api.example.com/mcp", {"Authorization": "Basic test"}) |
| 363 | + |
| 364 | + assert len(result) == 1 |
| 365 | + assert result[0].name == "test_tool" |
| 366 | + assert result[0].description == "Test description" |
| 367 | + assert result[0].input_schema == {"type": "object", "properties": {"id": {"type": "string"}}} |
| 368 | + |
| 369 | + def test_fetch_mcp_tools_with_pagination(self): |
| 370 | + """Test fetching tools with multiple pages.""" |
| 371 | + # First page |
| 372 | + mock_tool1 = MagicMock() |
| 373 | + mock_tool1.name = "tool_1" |
| 374 | + mock_tool1.description = "Tool 1" |
| 375 | + mock_tool1.inputSchema = {} |
| 376 | + |
| 377 | + mock_result1 = MagicMock() |
| 378 | + mock_result1.tools = [mock_tool1] |
| 379 | + mock_result1.nextCursor = "cursor_page_2" |
| 380 | + |
| 381 | + # Second page |
| 382 | + mock_tool2 = MagicMock() |
| 383 | + mock_tool2.name = "tool_2" |
| 384 | + mock_tool2.description = "Tool 2" |
| 385 | + mock_tool2.inputSchema = None # Test None inputSchema |
| 386 | + |
| 387 | + mock_result2 = MagicMock() |
| 388 | + mock_result2.tools = [mock_tool2] |
| 389 | + mock_result2.nextCursor = None |
| 390 | + |
| 391 | + # Create mock session with pagination |
| 392 | + mock_session = AsyncMock() |
| 393 | + mock_session.initialize = AsyncMock() |
| 394 | + mock_session.list_tools = AsyncMock(side_effect=[mock_result1, mock_result2]) |
| 395 | + mock_session.__aenter__ = AsyncMock(return_value=mock_session) |
| 396 | + mock_session.__aexit__ = AsyncMock(return_value=None) |
| 397 | + |
| 398 | + @asynccontextmanager |
| 399 | + async def mock_streamable_client(endpoint, headers): |
| 400 | + yield (MagicMock(), MagicMock(), MagicMock()) |
| 401 | + |
| 402 | + with ( |
| 403 | + patch( |
| 404 | + "mcp.client.streamable_http.streamablehttp_client", |
| 405 | + side_effect=mock_streamable_client, |
| 406 | + ), |
| 407 | + patch("mcp.client.session.ClientSession", return_value=mock_session), |
| 408 | + patch("mcp.types.Implementation", MagicMock()), |
| 409 | + ): |
| 410 | + result = _fetch_mcp_tools("https://api.example.com/mcp", {}) |
| 411 | + |
| 412 | + assert len(result) == 2 |
| 413 | + assert result[0].name == "tool_1" |
| 414 | + assert result[1].name == "tool_2" |
| 415 | + assert result[1].input_schema == {} # None should become empty dict |
| 416 | + assert mock_session.list_tools.call_count == 2 |
0 commit comments