Spaces:
Sleeping
Sleeping
File size: 2,154 Bytes
116524e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import pytest
import asyncio
from unittest.mock import MagicMock, patch
from ace.integrations.mcp.config import MCPServerConfig
from ace.integrations.mcp.registry import SessionRegistry
from ace.integrations.mcp.errors import SessionNotFoundError
@pytest.fixture
def config():
return MCPServerConfig(session_ttl_seconds=1)
@pytest.fixture
def registry(config):
return SessionRegistry(config)
@pytest.mark.asyncio
async def test_get_or_create(registry):
with patch("ace.integrations.mcp.registry.ACELiteLLM") as mock_runner_cls:
mock_runner_cls.from_model.return_value = MagicMock()
# Create
s1 = await registry.get_or_create("s1")
assert s1.session_id == "s1"
assert s1.runner is not None
mock_runner_cls.from_model.assert_called_once_with("gpt-4o-mini")
# Get existing
s1_again = await registry.get_or_create("s1")
assert s1 is s1_again
assert mock_runner_cls.from_model.call_count == 1
@pytest.mark.asyncio
async def test_get_existing(registry):
with patch("ace.integrations.mcp.registry.ACELiteLLM"):
s1 = await registry.get_or_create("s1")
s1_get = await registry.get("s1")
assert s1 is s1_get
@pytest.mark.asyncio
async def test_get_not_found(registry):
with pytest.raises(SessionNotFoundError):
await registry.get("nonexistent")
@pytest.mark.asyncio
async def test_sweep_expired(registry):
with patch("ace.integrations.mcp.registry.ACELiteLLM"):
s1 = await registry.get_or_create("s1")
# Should not expire immediately
await registry.get("s1")
# Wait for TTL to pass (config TTL is 1 sec)
await asyncio.sleep(1.1)
with pytest.raises(SessionNotFoundError):
await registry.get("s1")
@pytest.mark.asyncio
async def test_delete(registry):
with patch("ace.integrations.mcp.registry.ACELiteLLM"):
await registry.get_or_create("s1")
await registry.delete("s1")
with pytest.raises(SessionNotFoundError):
await registry.get("s1")
|