chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,446 @@
|
||||
"""
|
||||
Comprehensive test suite for Unity instance routing.
|
||||
|
||||
These tests validate that set_active_instance correctly routes subsequent
|
||||
tool calls to the intended Unity instance across ALL tool categories.
|
||||
|
||||
DESIGN: Single source of truth via middleware state:
|
||||
- set_active_instance tool stores instance per session in UnityInstanceMiddleware
|
||||
- Middleware injects instance into ctx.set_state() for each tool call
|
||||
- get_unity_instance_from_context() reads from ctx.get_state()
|
||||
- All tools (GameObject, Script, Asset, etc.) use get_unity_instance_from_context()
|
||||
"""
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, MagicMock, patch
|
||||
from fastmcp import Context
|
||||
|
||||
from core.config import config
|
||||
from transport.unity_instance_middleware import UnityInstanceMiddleware
|
||||
from services.tools import get_unity_instance_from_context
|
||||
from services.tools.set_active_instance import set_active_instance as set_active_instance_tool
|
||||
from transport.models import SessionList, SessionDetails
|
||||
|
||||
|
||||
def _make_stateful_ctx(session_id: str) -> Mock:
|
||||
"""Build a Context shim whose set/get/delete_state share one dict.
|
||||
|
||||
The middleware now defers persistence to FastMCP's session-scoped state
|
||||
store, so a useful unit-test mock must actually round-trip values rather
|
||||
than returning fresh Mocks. This helper does that minimally.
|
||||
"""
|
||||
state: dict[str, object] = {}
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = session_id
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k, v: state.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state.get(k))
|
||||
ctx.delete_state = AsyncMock(side_effect=lambda k: state.pop(k, None))
|
||||
return ctx
|
||||
|
||||
|
||||
class TestInstanceRoutingBasics:
|
||||
"""Test basic middleware functionality."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_stores_and_retrieves_instance(self):
|
||||
"""Middleware should store and retrieve instance per session."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = _make_stateful_ctx("test-session-1")
|
||||
|
||||
await middleware.set_active_instance(ctx, "TestProject@abc123")
|
||||
|
||||
assert await middleware.get_active_instance(ctx) == "TestProject@abc123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_isolates_sessions(self):
|
||||
"""Two MCP sessions must not see each other's active instance.
|
||||
|
||||
Regression test for #1023: the prior implementation keyed on the
|
||||
peer-supplied client_id and collapsed multiple clients onto a shared
|
||||
record. Each ctx in this test holds its own private state dict, so
|
||||
leakage would surface as a cross-read.
|
||||
"""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
|
||||
ctx1 = _make_stateful_ctx("session-1")
|
||||
ctx2 = _make_stateful_ctx("session-2")
|
||||
|
||||
await middleware.set_active_instance(ctx1, "Project1@aaa")
|
||||
await middleware.set_active_instance(ctx2, "Project2@bbb")
|
||||
|
||||
assert await middleware.get_active_instance(ctx1) == "Project1@aaa"
|
||||
assert await middleware.get_active_instance(ctx2) == "Project2@bbb"
|
||||
|
||||
|
||||
class TestInstanceRoutingIntegration:
|
||||
"""Test that instance routing works end-to-end for all tool categories."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_middleware_injects_state_into_context(self):
|
||||
"""Middleware on_call_tool should inject instance into ctx state.
|
||||
|
||||
After this PR the middleware writes two distinct keys: a persistence
|
||||
key (``mcpforunity.active_instance``) when ``set_active_instance`` is
|
||||
called, and a per-request injection key (``unity_instance``) inside
|
||||
``on_call_tool``. Tools downstream read ``unity_instance``.
|
||||
"""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
|
||||
ctx = _make_stateful_ctx("test-session")
|
||||
|
||||
middleware_ctx = Mock()
|
||||
middleware_ctx.fastmcp_context = ctx
|
||||
|
||||
await middleware.set_active_instance(ctx, "TestProject@abc123")
|
||||
|
||||
async def mock_call_next(ctx):
|
||||
return {"success": True}
|
||||
|
||||
await middleware.on_call_tool(middleware_ctx, mock_call_next)
|
||||
|
||||
# The per-request injection must have happened so tools downstream
|
||||
# can read it; the persistence write is verified by the round-trip
|
||||
# test in TestInstanceRoutingBasics above.
|
||||
ctx.set_state.assert_any_call("unity_instance", "TestProject@abc123")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unity_instance_from_context_checks_state(self):
|
||||
"""get_unity_instance_from_context must read from ctx.get_state()."""
|
||||
ctx = Mock(spec=Context)
|
||||
|
||||
# Set up state storage (only source of truth now)
|
||||
state_storage = {"unity_instance": "Project@state123"}
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
# Call and verify
|
||||
result = await get_unity_instance_from_context(ctx)
|
||||
|
||||
assert result == "Project@state123", \
|
||||
"get_unity_instance_from_context must read from ctx.get_state()!"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_unity_instance_returns_none_when_not_set(self):
|
||||
"""Should return None when no instance is set."""
|
||||
ctx = Mock(spec=Context)
|
||||
|
||||
# Empty state storage
|
||||
state_storage = {}
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
result = await get_unity_instance_from_context(ctx)
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestInstanceRoutingToolCategories:
|
||||
"""Test instance routing for each tool category."""
|
||||
|
||||
def _create_mock_context_with_instance(self, instance_id: str):
|
||||
"""Helper to create a mock context with instance set via middleware."""
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "test-session"
|
||||
|
||||
# Set up state storage (only source of truth)
|
||||
state_storage = {"unity_instance": instance_id}
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
|
||||
return ctx
|
||||
|
||||
|
||||
|
||||
class TestInstanceRoutingHTTP:
|
||||
"""Validate HTTP-specific routing helpers."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_active_instance_http_transport(self, monkeypatch):
|
||||
"""set_active_instance should enumerate PluginHub sessions under HTTP."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "http-session"
|
||||
state_storage = {}
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
monkeypatch.setattr(config, "transport_mode", "http")
|
||||
fake_sessions = SessionList(
|
||||
sessions={
|
||||
"sess-1": SessionDetails(
|
||||
project="Ramble",
|
||||
hash="8e29de57",
|
||||
unity_version="6000.2.10f1",
|
||||
connected_at="2025-11-21T03:30:03.682353+00:00",
|
||||
)
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.PluginHub.get_sessions",
|
||||
AsyncMock(return_value=fake_sessions),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.get_unity_instance_middleware",
|
||||
lambda: middleware,
|
||||
)
|
||||
|
||||
result = await set_active_instance_tool(ctx, "Ramble@8e29de57")
|
||||
|
||||
assert result["success"] is True
|
||||
assert await middleware.get_active_instance(ctx) == "Ramble@8e29de57"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_active_instance_http_hash_only(self, monkeypatch):
|
||||
"""Hash-only selection should resolve via PluginHub registry."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "http-session-2"
|
||||
state_storage = {}
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
monkeypatch.setattr(config, "transport_mode", "http")
|
||||
fake_sessions = SessionList(
|
||||
sessions={
|
||||
"sess-99": SessionDetails(
|
||||
project="UnityMCPTests",
|
||||
hash="cc8756d4",
|
||||
unity_version="2021.3.45f2",
|
||||
connected_at="2025-11-21T03:37:01.501022+00:00",
|
||||
)
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.PluginHub.get_sessions",
|
||||
AsyncMock(return_value=fake_sessions),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.get_unity_instance_middleware",
|
||||
lambda: middleware,
|
||||
)
|
||||
|
||||
result = await set_active_instance_tool(ctx, "UnityMCPTests@cc8756d4")
|
||||
|
||||
assert result["success"] is True
|
||||
assert await middleware.get_active_instance(ctx) == "UnityMCPTests@cc8756d4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_active_instance_http_hash_missing(self, monkeypatch):
|
||||
"""Unknown hashes should surface a clear error."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "http-session-3"
|
||||
|
||||
monkeypatch.setattr(config, "transport_mode", "http")
|
||||
fake_sessions = SessionList(sessions={})
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.PluginHub.get_sessions",
|
||||
AsyncMock(return_value=fake_sessions),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.get_unity_instance_middleware",
|
||||
lambda: middleware,
|
||||
)
|
||||
|
||||
result = await set_active_instance_tool(ctx, "Unknown@deadbeef")
|
||||
|
||||
assert result["success"] is False
|
||||
assert "No Unity instances" in result["error"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_active_instance_http_hash_ambiguous(self, monkeypatch):
|
||||
"""Ambiguous hash prefixes should mirror stdio error messaging."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "http-session-4"
|
||||
|
||||
monkeypatch.setattr(config, "transport_mode", "http")
|
||||
fake_sessions = SessionList(
|
||||
sessions={
|
||||
"sess-a": SessionDetails(project="ProjA", hash="abc12345", unity_version="2022", connected_at="now"),
|
||||
"sess-b": SessionDetails(project="ProjB", hash="abc98765", unity_version="2022", connected_at="now"),
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.PluginHub.get_sessions",
|
||||
AsyncMock(return_value=fake_sessions),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.tools.set_active_instance.get_unity_instance_middleware",
|
||||
lambda: middleware,
|
||||
)
|
||||
|
||||
result = await set_active_instance_tool(ctx, "abc")
|
||||
|
||||
assert result["success"] is False
|
||||
assert "Name@hash" in result["error"]
|
||||
|
||||
|
||||
class TestInstanceRoutingRaceConditions:
|
||||
"""Test for race conditions and timing issues."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rapid_instance_switching(self):
|
||||
"""Rapidly switching instances should not cause routing errors."""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "test-session"
|
||||
|
||||
state_storage = {}
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
instances = ["Project1@aaa", "Project2@bbb", "Project3@ccc"]
|
||||
|
||||
for instance in instances:
|
||||
await middleware.set_active_instance(ctx, instance)
|
||||
|
||||
# Create middleware context
|
||||
middleware_ctx = Mock()
|
||||
middleware_ctx.fastmcp_context = ctx
|
||||
|
||||
async def mock_call_next(ctx):
|
||||
return {"success": True}
|
||||
|
||||
# Execute middleware
|
||||
await middleware.on_call_tool(middleware_ctx, mock_call_next)
|
||||
|
||||
# Verify correct instance is set
|
||||
assert state_storage.get("unity_instance") == instance
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_then_immediate_create_script(self):
|
||||
"""Setting instance then immediately creating script should route correctly."""
|
||||
# This reproduces the bug: set_active_instance → create_script went to wrong instance
|
||||
|
||||
middleware = UnityInstanceMiddleware()
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "test-session"
|
||||
ctx.info = Mock()
|
||||
|
||||
state_storage = {}
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
ctx.request_context = None
|
||||
|
||||
# Set active instance
|
||||
await middleware.set_active_instance(ctx, "ramble@8e29de57")
|
||||
|
||||
# Simulate middleware intercepting create_script call
|
||||
middleware_ctx = Mock()
|
||||
middleware_ctx.fastmcp_context = ctx
|
||||
|
||||
async def mock_create_script_call(ctx):
|
||||
# This simulates what create_script does
|
||||
instance = await get_unity_instance_from_context(ctx)
|
||||
return {"success": True, "routed_to": instance}
|
||||
|
||||
# Inject state via middleware
|
||||
await middleware.on_call_tool(middleware_ctx, mock_create_script_call)
|
||||
|
||||
# Verify create_script would route to correct instance
|
||||
result = await mock_create_script_call(ctx)
|
||||
assert result["routed_to"] == "ramble@8e29de57", \
|
||||
"create_script must route to the instance set by set_active_instance"
|
||||
|
||||
|
||||
class TestInstanceRoutingSequentialOperations:
|
||||
"""Test the exact failure scenario from user report."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_four_script_creation_sequence(self):
|
||||
"""
|
||||
Reproduce the exact failure:
|
||||
1. set_active(ramble) → create_script1 → should go to ramble
|
||||
2. set_active(UnityMCPTests) → create_script2 → should go to UnityMCPTests
|
||||
3. set_active(ramble) → create_script3 → should go to ramble
|
||||
4. set_active(UnityMCPTests) → create_script4 → should go to UnityMCPTests
|
||||
|
||||
ACTUAL BEHAVIOR:
|
||||
- Script1 went to UnityMCPTests (WRONG)
|
||||
- Script2 went to ramble (WRONG)
|
||||
- Script3 went to ramble (CORRECT)
|
||||
- Script4 went to UnityMCPTests (CORRECT)
|
||||
"""
|
||||
middleware = UnityInstanceMiddleware()
|
||||
|
||||
# Track which instance each script was created in
|
||||
script_routes = {}
|
||||
|
||||
async def simulate_create_script(ctx, script_name, expected_instance):
|
||||
# Inject state via middleware
|
||||
middleware_ctx = Mock()
|
||||
middleware_ctx.fastmcp_context = ctx
|
||||
|
||||
async def mock_tool_call(middleware_ctx):
|
||||
# The middleware passes the middleware_ctx, we need the fastmcp_context
|
||||
tool_ctx = middleware_ctx.fastmcp_context
|
||||
instance = await get_unity_instance_from_context(tool_ctx)
|
||||
script_routes[script_name] = instance
|
||||
return {"success": True}
|
||||
|
||||
await middleware.on_call_tool(middleware_ctx, mock_tool_call)
|
||||
return expected_instance
|
||||
|
||||
# Session context
|
||||
ctx = Mock(spec=Context)
|
||||
ctx.session_id = "test-session"
|
||||
ctx.info = Mock()
|
||||
|
||||
state_storage = {}
|
||||
ctx.set_state = AsyncMock(side_effect=lambda k,
|
||||
v: state_storage.__setitem__(k, v))
|
||||
ctx.get_state = AsyncMock(side_effect=lambda k: state_storage.get(k))
|
||||
|
||||
# Execute sequence
|
||||
await middleware.set_active_instance(ctx, "ramble@8e29de57")
|
||||
expected1 = await simulate_create_script(ctx, "Script1", "ramble@8e29de57")
|
||||
|
||||
await middleware.set_active_instance(ctx, "UnityMCPTests@cc8756d4")
|
||||
expected2 = await simulate_create_script(ctx, "Script2", "UnityMCPTests@cc8756d4")
|
||||
|
||||
await middleware.set_active_instance(ctx, "ramble@8e29de57")
|
||||
expected3 = await simulate_create_script(ctx, "Script3", "ramble@8e29de57")
|
||||
|
||||
await middleware.set_active_instance(ctx, "UnityMCPTests@cc8756d4")
|
||||
expected4 = await simulate_create_script(ctx, "Script4", "UnityMCPTests@cc8756d4")
|
||||
|
||||
# Assertions - these will FAIL until the bug is fixed
|
||||
assert script_routes.get("Script1") == expected1, \
|
||||
f"Script1 should route to {expected1}, got {script_routes.get('Script1')}"
|
||||
assert script_routes.get("Script2") == expected2, \
|
||||
f"Script2 should route to {expected2}, got {script_routes.get('Script2')}"
|
||||
assert script_routes.get("Script3") == expected3, \
|
||||
f"Script3 should route to {expected3}, got {script_routes.get('Script3')}"
|
||||
assert script_routes.get("Script4") == expected4, \
|
||||
f"Script4 should route to {expected4}, got {script_routes.get('Script4')}"
|
||||
|
||||
|
||||
# Test regimen summary
|
||||
"""
|
||||
COMPREHENSIVE TEST REGIMEN FOR INSTANCE ROUTING
|
||||
|
||||
Prerequisites:
|
||||
- Two Unity instances running (e.g., ramble, UnityMCPTests)
|
||||
- MCP server connected to both instances
|
||||
|
||||
Test Categories:
|
||||
1. ✅ Middleware State Management (4 tests)
|
||||
2. ✅ Middleware Integration (2 tests)
|
||||
3. ✅ get_unity_instance_from_context (2 tests)
|
||||
4. ✅ Tool Category Coverage (11 categories)
|
||||
5. ✅ Race Conditions (2 tests)
|
||||
6. ✅ Sequential Operations (1 test - reproduces exact user bug)
|
||||
|
||||
Total: 21 tests
|
||||
|
||||
DESIGN:
|
||||
Single source of truth via middleware state:
|
||||
- set_active_instance stores instance per session in UnityInstanceMiddleware
|
||||
- Middleware injects instance into ctx.set_state() for each tool call
|
||||
- get_unity_instance_from_context() reads from ctx.get_state()
|
||||
- All tools use get_unity_instance_from_context()
|
||||
|
||||
This ensures consistent routing across ALL tool categories (Script, GameObject, Asset, etc.)
|
||||
"""
|
||||
Reference in New Issue
Block a user