test_domain_reload_resilience.py•7.63 kB
"""
Integration test for domain reload resilience.
Tests that the MCP server can handle rapid-fire requests during Unity domain reloads
by waiting for the plugin to reconnect instead of failing immediately.
"""
import asyncio
import pytest
from unittest.mock import AsyncMock, patch
from datetime import datetime
from .test_helpers import DummyContext
@pytest.mark.asyncio
async def test_plugin_hub_waits_for_reconnection_during_reload():
"""Test that PluginHub._resolve_session_id waits for plugin reconnection."""
# Import after conftest stubs are set up
from transport.plugin_hub import PluginHub
from transport.plugin_registry import PluginRegistry, PluginSession
# Create a mock registry
mock_registry = AsyncMock(spec=PluginRegistry)
# Simulate plugin reconnection sequence:
# First 2 calls: no sessions (plugin disconnected)
# Third call: session appears (plugin reconnected)
call_count = [0]
async def mock_list_sessions():
call_count[0] += 1
if call_count[0] <= 2:
# Plugin not yet reconnected
return {}
else:
# Plugin reconnected
now = datetime.now()
session = PluginSession(
session_id="test-session-123",
project_name="TestProject",
project_hash="abc123",
unity_version="2022.3.0f1",
registered_at=now,
connected_at=now
)
return {"test-session-123": session}
mock_registry.list_sessions = mock_list_sessions
# Configure PluginHub with our mock while preserving the original state
original_registry = PluginHub._registry
original_lock = PluginHub._lock
PluginHub._registry = mock_registry
PluginHub._lock = asyncio.Lock()
try:
# Call _resolve_session_id when no session is available
# It should wait and retry until the session appears
session_id = await PluginHub._resolve_session_id(unity_instance=None)
# Should have retried and eventually found the session
assert session_id == "test-session-123"
assert call_count[0] >= 3 # Should have tried at least 3 times
finally:
# Clean up: restore original PluginHub state
PluginHub._registry = original_registry
PluginHub._lock = original_lock
@pytest.mark.asyncio
async def test_plugin_hub_fails_after_timeout():
"""Test that PluginHub._resolve_session_id eventually times out if plugin never reconnects."""
from transport.plugin_hub import PluginHub
from transport.plugin_registry import PluginRegistry
# Create a mock registry that never returns sessions
mock_registry = AsyncMock(spec=PluginRegistry)
async def mock_list_sessions():
return {} # Never returns sessions
mock_registry.list_sessions = mock_list_sessions
# Configure PluginHub with our mock while preserving the original state
original_registry = PluginHub._registry
original_lock = PluginHub._lock
PluginHub._registry = mock_registry
PluginHub._lock = asyncio.Lock()
# Temporarily override config for a short timeout
with patch('transport.plugin_hub.config') as mock_config:
mock_config.reload_max_retries = 3 # Only 3 retries
mock_config.reload_retry_ms = 10 # 10ms between retries
try:
# Should raise RuntimeError after timeout
with pytest.raises(RuntimeError, match="No Unity plugins are currently connected"):
await PluginHub._resolve_session_id(unity_instance=None)
finally:
# Clean up: restore original PluginHub state
PluginHub._registry = original_registry
PluginHub._lock = original_lock
@pytest.mark.asyncio
async def test_read_console_during_simulated_reload(monkeypatch):
"""
Simulate the stress test: create script (triggers reload) + rapid read_console calls.
This test simulates what happens when:
1. A script is created (triggering domain reload)
2. Multiple read_console calls are made immediately
3. The plugin disconnects and reconnects during those calls
"""
# Setup tools
from services.tools.read_console import read_console
call_count = [0]
async def fake_send_command(*args, **kwargs):
"""Simulate successful command execution."""
call_count[0] += 1
return {
"success": True,
"message": f"Retrieved {call_count[0]} log entries.",
"data": ["<b><color=#2EA3FF>MCP-FOR-UNITY</color></b>: Auto-discovered 10 tools"]
}
# Patch the async_send_command_with_retry directly
import services.tools.read_console
monkeypatch.setattr(
services.tools.read_console,
"async_send_command_with_retry",
fake_send_command
)
# Run multiple read_console calls rapidly (simulating the stress test)
results = []
for i in range(5):
result = await read_console(
ctx=DummyContext(),
action="get",
types=["all"],
count=50,
format="plain",
include_stacktrace=False
)
results.append(result)
# All calls should succeed
assert len(results) == 5
for i, result in enumerate(results):
assert result["success"] is True, f"Call {i+1} failed with result: {result}"
assert "data" in result
# At least 5 calls should have been made
assert call_count[0] == 5
@pytest.mark.asyncio
async def test_plugin_hub_respects_unity_instance_preference():
"""Test that _resolve_session_id prefers a specific Unity instance if requested."""
from transport.plugin_hub import PluginHub
from transport.plugin_registry import PluginRegistry, PluginSession
# Create a mock registry with two sessions
mock_registry = AsyncMock(spec=PluginRegistry)
now = datetime.now()
session1 = PluginSession(
session_id="session-1",
project_name="Project1",
project_hash="hash1",
unity_version="2022.3.0f1",
registered_at=now,
connected_at=now
)
session2 = PluginSession(
session_id="session-2",
project_name="Project2",
project_hash="hash2",
unity_version="2022.3.0f1",
registered_at=now,
connected_at=now
)
async def mock_list_sessions():
return {
"session-1": session1,
"session-2": session2
}
async def mock_get_session_by_hash(project_hash):
if project_hash == "hash2":
return "session-2"
return None
mock_registry.list_sessions = mock_list_sessions
mock_registry.get_session_id_by_hash = mock_get_session_by_hash
# Configure PluginHub with our mock while preserving the original state
original_registry = PluginHub._registry
original_lock = PluginHub._lock
PluginHub._registry = mock_registry
PluginHub._lock = asyncio.Lock()
try:
# Request specific Unity instance
session_id = await PluginHub._resolve_session_id(unity_instance="hash2")
# Should return the requested instance
assert session_id == "session-2"
# Request default (no specific instance)
with pytest.raises(RuntimeError) as exc:
await PluginHub._resolve_session_id(unity_instance=None)
assert "Multiple Unity instances are connected" in str(exc.value)
finally:
# Clean up: restore original PluginHub state
PluginHub._registry = original_registry
PluginHub._lock = original_lock