OpenHands/tests/unit/test_mcp_client_timeout.py
Rohit Malhotra a348840534
[Feat]: support streamable http mcp (#8864)
Co-authored-by: openhands <openhands@all-hands.dev>
2025-06-03 17:06:44 +00:00

50 lines
1.7 KiB
Python

import asyncio
from contextlib import asynccontextmanager
from unittest import mock
import pytest
from openhands.mcp.client import MCPClient
@pytest.mark.asyncio
async def test_connect_sse_timeout():
"""Test that connect_sse properly times out when server_url is invalid."""
client = MCPClient()
# Create a mock async context manager that simulates a timeout
@asynccontextmanager
async def mock_slow_context(*args, **kwargs):
# This will hang for longer than our timeout
await asyncio.sleep(10.0)
yield (mock.AsyncMock(), mock.AsyncMock())
# Patch the sse_client function to return our slow context manager
with mock.patch(
'openhands.mcp.client.sse_client', return_value=mock_slow_context()
):
# Test with a very short timeout
with pytest.raises(asyncio.TimeoutError):
await client.connect_sse('http://example.com', timeout=0.1)
@pytest.mark.asyncio
async def test_connect_streamable_http_timeout():
"""Test that connect_streamable_http properly times out when server_url is invalid."""
client = MCPClient()
# Create a mock async context manager that simulates a timeout
@asynccontextmanager
async def mock_slow_context(*args, **kwargs):
# This will hang for longer than our timeout
await asyncio.sleep(10.0)
yield (mock.AsyncMock(), mock.AsyncMock(), mock.AsyncMock())
# Patch the streamablehttp_client function to return our slow context manager
with mock.patch(
'openhands.mcp.client.streamablehttp_client', return_value=mock_slow_context()
):
# Test with a very short timeout
with pytest.raises(asyncio.TimeoutError):
await client.connect_shttp('http://example.com', timeout=0.1)