"""Tests for the behavioral fixes in fix/mcp-server-bugs. Covers: - _set_client invalidates shared registry on client change - login() session isolation (only closes contextvar client) - create_flow_from_spec raises on validation failure - validate_flow filters builds by job_id - update_flow_from_spec rollback on failure """ import copy from unittest.mock import AsyncMock, MagicMock, patch import pytest # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_mock_client(**overrides): client = AsyncMock() client.post_event = AsyncMock() client.get = AsyncMock(return_value={}) client.patch = AsyncMock(return_value={}) client.post = AsyncMock(return_value={}) client.close = AsyncMock() client.login = AsyncMock() client.server_url = "http://localhost:7860" for k, v in overrides.items(): setattr(client, k, v) return client def _mock_flow(flow_id="flow-123", nodes=None, edges=None): return { "id": flow_id, "name": "Test Flow", "data": { "nodes": nodes or [], "edges": edges or [], "viewport": {"x": 0, "y": 0, "zoom": 1}, }, } MOCK_REGISTRY = { "ChatInput": { "display_name": "Chat Input", "template": {}, "output_types": ["Message"], "outputs": [{"name": "message", "types": ["Message"]}], }, } # --------------------------------------------------------------------------- # _set_client: invalidates shared registry when client changes # --------------------------------------------------------------------------- class TestSetClientRegistryInvalidation: def test_invalidates_registry_on_new_client(self): import lfx.mcp.server as srv old_client = MagicMock() new_client = MagicMock() srv._shared_client = old_client srv._shared_registry = {"some": "registry"} srv._set_client(new_client) assert srv._shared_registry is None assert srv._shared_client is new_client def test_preserves_registry_when_same_client(self): import lfx.mcp.server as srv client = MagicMock() registry = {"some": "registry"} srv._shared_client = client srv._shared_registry = registry srv._set_client(client) assert srv._shared_registry is registry assert srv._shared_client is client # --------------------------------------------------------------------------- # login(): session isolation -- only closes the contextvar client # --------------------------------------------------------------------------- class TestLoginSessionIsolation: async def test_login_closes_only_session_client(self): """login() should close the contextvar client, not the shared one.""" import lfx.mcp.server as srv session_client = _make_mock_client() shared_client = _make_mock_client() # Set up: session has its own client, shared has a different one srv._client_var.set(session_client) srv._shared_client = shared_client with patch.object(srv, "LangflowClient", return_value=_make_mock_client()): await srv.login("user", "pass", "http://localhost:7860") session_client.close.assert_called_once() shared_client.close.assert_not_called() async def test_login_no_close_when_no_session_client(self): """login() should not close anything when no contextvar client exists.""" import lfx.mcp.server as srv shared_client = _make_mock_client() srv._client_var.set(None) srv._shared_client = shared_client new_client = _make_mock_client() with patch.object(srv, "LangflowClient", return_value=new_client): await srv.login("user", "pass", "http://localhost:7860") shared_client.close.assert_not_called() async def test_login_clears_session_registry_not_shared(self): """login() should clear _registry_var but leave _shared_registry alone.""" import lfx.mcp.server as srv srv._client_var.set(None) srv._shared_client = _make_mock_client() srv._shared_registry = {"cached": "registry"} srv._registry_var.set({"session": "registry"}) new_client = _make_mock_client() with patch.object(srv, "LangflowClient", return_value=new_client): await srv.login("user", "pass", "http://localhost:7860") assert srv._registry_var.get() is None # _shared_registry may be invalidated by _set_client if client changes, # but login itself should not explicitly clear it # --------------------------------------------------------------------------- # create_flow_from_spec: raises RuntimeError when validation fails # --------------------------------------------------------------------------- class TestCreateFlowFromSpecValidation: async def test_raises_on_validation_failure(self): from lfx.mcp.server import create_flow_from_spec parsed = { "name": "T", "description": "", "nodes": [{"id": "A", "type": "ChatInput"}], "edges": [], "config": {}, } created = {"id": "flow-new", "name": "T", "description": ""} mock_client = _make_mock_client() validation_result = { "valid": False, "errors": [{"component_id": "X-1", "error": "Missing API key"}], } with ( patch("lfx.mcp.server._get_client", return_value=mock_client), patch("lfx.mcp.server._get_registry", new_callable=AsyncMock, return_value=MOCK_REGISTRY), patch("lfx.mcp.server.parse_flow_spec", return_value=parsed), patch("lfx.mcp.server.validate_spec_references"), patch("lfx.mcp.server.create_flow", new_callable=AsyncMock, return_value=created), patch("lfx.mcp.server.add_component", new_callable=AsyncMock, return_value={"id": "X-1"}), patch("lfx.mcp.server.validate_flow", new_callable=AsyncMock, return_value=validation_result), patch("lfx.mcp.server.delete_flow", new_callable=AsyncMock), pytest.raises(RuntimeError, match=r"Flow validation failed.*Missing API key"), ): await create_flow_from_spec("nodes:\n A: ChatInput") async def test_skips_validation_when_disabled(self): from lfx.mcp.server import create_flow_from_spec parsed = { "name": "T", "description": "", "nodes": [{"id": "A", "type": "ChatInput"}], "edges": [], "config": {}, } created = {"id": "flow-new", "name": "T", "description": ""} mock_client = _make_mock_client() with ( patch("lfx.mcp.server._get_client", return_value=mock_client), patch("lfx.mcp.server._get_registry", new_callable=AsyncMock, return_value=MOCK_REGISTRY), patch("lfx.mcp.server.parse_flow_spec", return_value=parsed), patch("lfx.mcp.server.validate_spec_references"), patch("lfx.mcp.server.create_flow", new_callable=AsyncMock, return_value=created), patch("lfx.mcp.server.add_component", new_callable=AsyncMock, return_value={"id": "X-1"}), patch("lfx.mcp.server.validate_flow", new_callable=AsyncMock) as mock_validate, patch("lfx.mcp.server.get_flow_info", new_callable=AsyncMock, return_value={"id": "flow-new"}), ): await create_flow_from_spec("nodes:\n A: ChatInput", validate=False) mock_validate.assert_not_called() # --------------------------------------------------------------------------- # update_flow_from_spec: rollback on failure # --------------------------------------------------------------------------- class TestUpdateFlowFromSpecRollback: async def test_restores_original_data_on_failure(self): from lfx.mcp.server import update_flow_from_spec mock_client = _make_mock_client() original_nodes = [{"id": "n1", "data": {"id": "orig-comp"}}] flow = _mock_flow(flow_id="flow-123", nodes=original_nodes) original_data = copy.deepcopy(flow["data"]) parsed = { "name": "T", "description": "", "nodes": [{"id": "A", "type": "ChatInput"}], "edges": [], "config": {}, } with ( patch("lfx.mcp.server._get_client", return_value=mock_client), patch("lfx.mcp.server._get_flow", new_callable=AsyncMock, return_value=flow), patch("lfx.mcp.server._get_registry", new_callable=AsyncMock, return_value=MOCK_REGISTRY), patch("lfx.mcp.server.parse_flow_spec", return_value=parsed), patch("lfx.mcp.server.validate_spec_references"), patch("lfx.mcp.server.add_component", new_callable=AsyncMock, side_effect=RuntimeError("Server error")), pytest.raises(RuntimeError, match="Server error"), ): await update_flow_from_spec("flow-123", "nodes:\n A: ChatInput") # Verify rollback: the last patch call should restore original_data patch_calls = mock_client.patch.call_args_list last_patch = patch_calls[-1] assert last_patch[1]["json_data"]["data"] == original_data async def test_emits_settle_event_on_rollback(self): from lfx.mcp.server import update_flow_from_spec mock_client = _make_mock_client() flow = _mock_flow(flow_id="flow-123", nodes=[{"id": "n1", "data": {"id": "orig-comp"}}]) parsed = { "name": "T", "description": "", "nodes": [{"id": "A", "type": "ChatInput"}], "edges": [], "config": {}, } with ( patch("lfx.mcp.server._get_client", return_value=mock_client), patch("lfx.mcp.server._get_flow", new_callable=AsyncMock, return_value=flow), patch("lfx.mcp.server._get_registry", new_callable=AsyncMock, return_value=MOCK_REGISTRY), patch("lfx.mcp.server.parse_flow_spec", return_value=parsed), patch("lfx.mcp.server.validate_spec_references"), patch("lfx.mcp.server.add_component", new_callable=AsyncMock, side_effect=RuntimeError("fail")), pytest.raises(RuntimeError, match="fail"), ): await update_flow_from_spec("flow-123", "nodes:\n A: ChatInput") # Should have emitted flow_settled before restoring mock_client.post_event.assert_called_once() args = mock_client.post_event.call_args[0] assert args[0] == "flow-123" assert args[1] == "flow_settled" async def test_success_does_not_rollback(self): from lfx.mcp.server import update_flow_from_spec mock_client = _make_mock_client() flow = _mock_flow(flow_id="flow-123") parsed = { "name": "Updated", "description": "", "nodes": [{"id": "A", "type": "ChatInput"}], "edges": [], "config": {}, } with ( patch("lfx.mcp.server._get_client", return_value=mock_client), patch("lfx.mcp.server._get_flow", new_callable=AsyncMock, return_value=flow), patch("lfx.mcp.server._get_registry", new_callable=AsyncMock, return_value=MOCK_REGISTRY), patch("lfx.mcp.server.parse_flow_spec", return_value=parsed), patch("lfx.mcp.server.validate_spec_references"), patch("lfx.mcp.server.add_component", new_callable=AsyncMock, return_value={"id": "X-1"}), patch("lfx.mcp.server._create_prompt_template_vars", new_callable=AsyncMock), patch("lfx.mcp.server.get_flow_info", new_callable=AsyncMock, return_value={"id": "flow-123"}), ): result = await update_flow_from_spec("flow-123", "nodes:\n A: ChatInput") assert result["id"] == "flow-123" assert result["node_id_map"] == {"A": "X-1"} # Only 2 patch calls: initial clear + no rollback # First patch clears nodes/edges, no second rollback patch patch_calls = [c for c in mock_client.patch.call_args_list if "data" in (c[1].get("json_data") or {})] assert len(patch_calls) == 1 # Only the initial clear