import asyncio from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import pytest from fastapi import HTTPException, status from httpx import AsyncClient from langflow.api.v1.mcp_projects import ( get_project_mcp_server, get_project_sse, init_mcp_servers, project_mcp_servers, project_sse_transports, ) from langflow.services.auth.utils import create_user_longterm_token, get_password_hash from langflow.services.database.models.flow import Flow from langflow.services.database.models.folder import Folder from langflow.services.database.models.user.model import User from langflow.services.deps import get_db_service, get_settings_service, session_scope from mcp.server.sse import SseServerTransport from sqlmodel import select # Mark all tests in this module as asyncio pytestmark = pytest.mark.asyncio @pytest.fixture def mock_project(active_user): """Fixture to provide a mock project linked to the active user.""" return Folder(id=uuid4(), name="Test Project", user_id=active_user.id) @pytest.fixture def mock_flow(active_user, mock_project): """Fixture to provide a mock flow linked to the active user and project.""" return Flow( id=uuid4(), name="Test Flow", description="Test Description", mcp_enabled=True, action_name="test_action", action_description="Test Action Description", folder_id=mock_project.id, user_id=active_user.id, ) @pytest.fixture def mock_project_mcp_server(): with patch("langflow.api.v1.mcp_projects.ProjectMCPServer") as mock: server_instance = MagicMock() server_instance.server = MagicMock() server_instance.server.name = "test-server" server_instance.server.run = AsyncMock() server_instance.server.create_initialization_options = MagicMock() mock.return_value = server_instance yield server_instance class AsyncContextManagerMock: """Mock class that implements async context manager protocol.""" async def __aenter__(self): return (MagicMock(), MagicMock()) async def __aexit__(self, exc_type, exc_val, exc_tb): # No teardown required for this mock context manager in tests pass @pytest.fixture def mock_sse_transport(): with patch("langflow.api.v1.mcp_projects.SseServerTransport") as mock: transport_instance = MagicMock() # Create an async context manager for connect_sse connect_sse_mock = AsyncContextManagerMock() transport_instance.connect_sse = MagicMock(return_value=connect_sse_mock) transport_instance.handle_post_message = AsyncMock() mock.return_value = transport_instance yield transport_instance @pytest.fixture(autouse=True) def mock_current_user_ctx(active_user): with patch("langflow.api.v1.mcp_projects.current_user_ctx") as mock: mock.get.return_value = active_user mock.set = MagicMock(return_value="dummy_token") mock.reset = MagicMock() yield mock @pytest.fixture(autouse=True) def mock_current_project_ctx(mock_project): with patch("langflow.api.v1.mcp_projects.current_project_ctx") as mock: mock.get.return_value = mock_project.id mock.set = MagicMock(return_value="dummy_token") mock.reset = MagicMock() yield mock @pytest.fixture async def other_test_user(): """Fixture for creating another test user.""" user_id = uuid4() async with session_scope() as session: user = User( id=user_id, username="other_test_user", password=get_password_hash("testpassword"), is_active=True, is_superuser=False, ) session.add(user) await session.commit() await session.refresh(user) yield user # Clean up async with session_scope() as session: user = await session.get(User, user_id) if user: await session.delete(user) await session.commit() @pytest.fixture async def other_test_project(other_test_user): """Fixture for creating a project for another test user.""" project_id = uuid4() async with session_scope() as session: project = Folder(id=project_id, name="Other Test Project", user_id=other_test_user.id) session.add(project) await session.commit() await session.refresh(project) yield project # Clean up async with session_scope() as session: project = await session.get(Folder, project_id) if project: await session.delete(project) await session.commit() @pytest.fixture(autouse=True) def disable_mcp_composer_by_default(): """Auto-fixture to disable MCP Composer for all tests by default.""" with patch("langflow.api.v1.mcp_projects.get_settings_service") as mock_get_settings: from langflow.services.deps import get_settings_service real_service = get_settings_service() # Create a mock that returns False for mcp_composer_enabled but delegates everything else mock_service = MagicMock() mock_service.settings = MagicMock() mock_service.settings.mcp_composer_enabled = False # Copy any other settings that might be needed mock_service.auth_settings = real_service.auth_settings mock_get_settings.return_value = mock_service yield @pytest.fixture def enable_mcp_composer(): """Fixture to explicitly enable MCP Composer for specific tests.""" with patch("langflow.api.v1.mcp_projects.get_settings_service") as mock_get_settings: from langflow.services.deps import get_settings_service real_service = get_settings_service() mock_service = MagicMock() mock_service.settings = MagicMock() mock_service.settings.mcp_composer_enabled = True # Copy any other settings that might be needed mock_service.auth_settings = real_service.auth_settings mock_get_settings.return_value = mock_service yield True async def test_handle_project_messages_success( client: AsyncClient, user_test_project, mock_sse_transport, logged_in_headers ): """Test successful handling of project messages.""" response = await client.post( f"api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, json={"type": "test", "content": "message"}, ) assert response.status_code == status.HTTP_200_OK mock_sse_transport.handle_post_message.assert_called_once() async def test_update_project_mcp_settings_invalid_json(client: AsyncClient, user_test_project, logged_in_headers): """Test updating MCP settings with invalid JSON.""" response = await client.patch( f"api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, json="invalid" ) assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY @pytest.fixture async def test_flow_for_update(active_user, user_test_project): """Fixture to provide a real flow for testing MCP settings updates.""" flow_id = uuid4() flow_data = { "id": flow_id, "name": "Test Flow For Update", "description": "Test flow that will be updated", "mcp_enabled": True, "action_name": "original_action", "action_description": "Original description", "folder_id": user_test_project.id, "user_id": active_user.id, } # Create the flow in the database async with session_scope() as session: flow = Flow(**flow_data) session.add(flow) await session.commit() await session.refresh(flow) yield flow # Clean up async with session_scope() as session: # Get the flow from the database flow = await session.get(Flow, flow_id) if flow: await session.delete(flow) await session.commit() async def test_update_project_mcp_settings_success( client: AsyncClient, user_test_project, test_flow_for_update, logged_in_headers ): """Test successful update of MCP settings using real database.""" # Create settings for updating the flow json_payload = { "settings": [ { "id": str(test_flow_for_update.id), "action_name": "updated_action", "action_description": "Updated description", "mcp_enabled": False, "name": test_flow_for_update.name, "description": test_flow_for_update.description, } ], "auth_settings": { "auth_type": "none", "api_key": None, "iam_endpoint": None, "username": None, "password": None, "bearer_token": None, }, } # Make the real PATCH request response = await client.patch( f"api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, json=json_payload ) # Assert response assert response.status_code == 200 assert "Updated MCP settings for 1 flows" in response.json()["message"] # Verify the flow was actually updated in the database async with session_scope() as session: updated_flow = await session.get(Flow, test_flow_for_update.id) assert updated_flow is not None assert updated_flow.action_name == "updated_action" assert updated_flow.action_description == "Updated description" assert updated_flow.mcp_enabled is False async def test_update_project_mcp_settings_invalid_project(client: AsyncClient, logged_in_headers): """Test accessing an invalid project ID.""" # We're using the GET endpoint since it works correctly and tests the same security constraints # Generate a random UUID that doesn't exist in the database nonexistent_project_id = uuid4() # Try to access the project response = await client.get(f"api/v1/mcp/project/{nonexistent_project_id}/sse", headers=logged_in_headers) # Verify the response assert response.status_code == 404 assert response.json()["detail"] == "Project not found" async def test_update_project_mcp_settings_other_user_project( client: AsyncClient, other_test_project, logged_in_headers ): """Test accessing a project belonging to another user.""" # We're using the GET endpoint since it works correctly and tests the same security constraints # This test disables MCP Composer to test JWT token-based access control # Try to access the other user's project using active_user's credentials response = await client.get(f"api/v1/mcp/project/{other_test_project.id}/sse", headers=logged_in_headers) # Verify the response assert response.status_code == 404 assert response.json()["detail"] == "Project not found" async def test_update_project_mcp_settings_other_user_project_with_composer( client: AsyncClient, other_test_project, logged_in_headers, enable_mcp_composer ): """Test accessing a project belonging to another user when MCP Composer is enabled.""" # When MCP Composer is enabled, JWT tokens are not accepted for MCP endpoints assert enable_mcp_composer # Fixture ensures MCP Composer is enabled # Try to access the other user's project using active_user's JWT credentials response = await client.get(f"api/v1/mcp/project/{other_test_project.id}/sse", headers=logged_in_headers) # Verify the response - should get 401 because JWT tokens aren't accepted assert response.status_code == 401 assert "API key required" in response.json()["detail"] async def test_update_project_mcp_settings_empty_settings(client: AsyncClient, user_test_project, logged_in_headers): """Test updating MCP settings with empty settings list.""" # Use real database objects instead of mocks to avoid the coroutine issue # Empty settings list json_payload = { "settings": [], "auth_settings": { "auth_type": "none", "api_key": None, "iam_endpoint": None, "username": None, "password": None, "bearer_token": None, }, } # Make the request to the actual endpoint response = await client.patch( f"api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, json=json_payload ) # Verify response - the real endpoint should handle empty settings correctly assert response.status_code == 200 assert "Updated MCP settings for 0 flows" in response.json()["message"] async def test_user_can_only_access_own_projects(client: AsyncClient, other_test_project, logged_in_headers): """Test that a user can only access their own projects.""" # Try to access the other user's project using first user's credentials response = await client.get(f"api/v1/mcp/project/{other_test_project.id}/sse", headers=logged_in_headers) # Should fail with 404 as first user cannot see second user's project assert response.status_code == 404 assert response.json()["detail"] == "Project not found" async def test_user_data_isolation_with_real_db( client: AsyncClient, logged_in_headers, other_test_user, other_test_project ): """Test that users can only access their own MCP projects using a real database session.""" # Create a flow for the other test user in their project second_flow_id = uuid4() # Use real database session just for flow creation and cleanup async with session_scope() as session: # Create a flow in the other user's project second_flow = Flow( id=second_flow_id, name="Second User Flow", description="This flow belongs to the second user", mcp_enabled=True, action_name="second_user_action", action_description="Second user action description", folder_id=other_test_project.id, user_id=other_test_user.id, ) # Add flow to database session.add(second_flow) await session.commit() try: # Test that first user can't see the project response = await client.get(f"api/v1/mcp/project/{other_test_project.id}/sse", headers=logged_in_headers) # Should fail with 404 assert response.status_code == 404 assert response.json()["detail"] == "Project not found" # First user attempts to update second user's flow settings # Note: We're not testing the PATCH endpoint because it has the coroutine error # Instead, verify permissions via the GET endpoint finally: # Clean up flow async with session_scope() as session: second_flow = await session.get(Flow, second_flow_id) if second_flow: await session.delete(second_flow) await session.commit() @pytest.fixture async def user_test_project(active_user): """Fixture for creating a project for the active user.""" project_id = uuid4() async with session_scope() as session: project = Folder(id=project_id, name="User Test Project", user_id=active_user.id) session.add(project) await session.commit() await session.refresh(project) yield project # Clean up async with session_scope() as session: project = await session.get(Folder, project_id) if project: await session.delete(project) await session.commit() @pytest.fixture async def user_test_flow(active_user, user_test_project): """Fixture for creating a flow for the active user.""" flow_id = uuid4() async with session_scope() as session: flow = Flow( id=flow_id, name="User Test Flow", description="This flow belongs to the active user", mcp_enabled=True, action_name="user_action", action_description="User action description", folder_id=user_test_project.id, user_id=active_user.id, ) session.add(flow) await session.commit() await session.refresh(flow) yield flow # Clean up async with session_scope() as session: flow = await session.get(Flow, flow_id) if flow: await session.delete(flow) await session.commit() async def test_user_can_update_own_flow_mcp_settings( client: AsyncClient, logged_in_headers, user_test_project, user_test_flow ): """Test that a user can update MCP settings for their own flows using real database.""" # User attempts to update their own flow settings json_payload = { "settings": [ { "id": str(user_test_flow.id), "action_name": "updated_user_action", "action_description": "Updated user action description", "mcp_enabled": False, "name": "User Test Flow", "description": "This flow belongs to the active user", } ], "auth_settings": { "auth_type": "none", "api_key": None, "iam_endpoint": None, "username": None, "password": None, "bearer_token": None, }, } # Make the PATCH request to update settings response = await client.patch( f"api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, json=json_payload ) # Should succeed as the user owns this project and flow assert response.status_code == 200 assert "Updated MCP settings for 1 flows" in response.json()["message"] # Verify the flow was actually updated in the database async with session_scope() as session: updated_flow = await session.get(Flow, user_test_flow.id) assert updated_flow is not None assert updated_flow.action_name == "updated_user_action" assert updated_flow.action_description == "Updated user action description" assert updated_flow.mcp_enabled is False async def test_update_project_auth_settings_encryption( client: AsyncClient, user_test_project, test_flow_for_update, logged_in_headers ): """Test that sensitive auth_settings fields are encrypted when stored.""" # Create settings with sensitive data json_payload = { "settings": [ { "id": str(test_flow_for_update.id), "action_name": "test_action", "action_description": "Test description", "mcp_enabled": True, "name": test_flow_for_update.name, "description": test_flow_for_update.description, } ], "auth_settings": { "auth_type": "oauth", "oauth_host": "localhost", "oauth_port": "3000", "oauth_server_url": "http://localhost:3000", "oauth_callback_path": "/callback", "oauth_client_id": "test-client-id", "oauth_client_secret": "test-oauth-secret-value-456", "oauth_auth_url": "https://oauth.example.com/auth", "oauth_token_url": "https://oauth.example.com/token", "oauth_mcp_scope": "read write", "oauth_provider_scope": "user:email", }, } # Send the update request response = await client.patch( f"/api/v1/mcp/project/{user_test_project.id}", json=json_payload, headers=logged_in_headers, ) assert response.status_code == 200 # Verify the sensitive data is encrypted in the database async with session_scope() as session: updated_project = await session.get(Folder, user_test_project.id) assert updated_project is not None assert updated_project.auth_settings is not None # Check that sensitive field is encrypted (not plaintext) stored_value = updated_project.auth_settings.get("oauth_client_secret") assert stored_value is not None assert stored_value != "test-oauth-secret-value-456" # Should be encrypted # The encrypted value should be a base64-like string (Fernet token) assert len(stored_value) > 50 # Encrypted values are longer # Now test that the GET endpoint returns the data (SecretStr will be masked) response = await client.get( f"/api/v1/mcp/project/{user_test_project.id}", headers=logged_in_headers, ) assert response.status_code == 200 data = response.json() # SecretStr fields are masked in the response for security assert data["auth_settings"]["oauth_client_secret"] == "**********" # noqa: S105 assert data["auth_settings"]["oauth_client_id"] == "test-client-id" assert data["auth_settings"]["auth_type"] == "oauth" # Verify that decryption is working by checking the actual decrypted value in the backend from langflow.services.auth.mcp_encryption import decrypt_auth_settings async with session_scope() as session: project = await session.get(Folder, user_test_project.id) decrypted_settings = decrypt_auth_settings(project.auth_settings) assert decrypted_settings["oauth_client_secret"] == "test-oauth-secret-value-456" # noqa: S105 async def test_project_sse_creation(user_test_project): """Test that SSE transport and MCP server are correctly created for a project.""" # Test getting an SSE transport for the first time project_id = user_test_project.id project_id_str = str(project_id) # Ensure there's no SSE transport for this project yet if project_id_str in project_sse_transports: del project_sse_transports[project_id_str] # Get an SSE transport sse_transport = get_project_sse(project_id) # Verify the transport was created correctly assert project_id_str in project_sse_transports assert sse_transport is project_sse_transports[project_id_str] assert isinstance(sse_transport, SseServerTransport) # Test getting an MCP server for the first time if project_id_str in project_mcp_servers: del project_mcp_servers[project_id_str] # Get an MCP server mcp_server = get_project_mcp_server(project_id) # Verify the server was created correctly assert project_id_str in project_mcp_servers assert mcp_server is project_mcp_servers[project_id_str] assert mcp_server.project_id == project_id assert mcp_server.server.name == f"langflow-mcp-project-{project_id}" # Test that getting the same SSE transport and MCP server again returns the cached instances sse_transport2 = get_project_sse(project_id) mcp_server2 = get_project_mcp_server(project_id) assert sse_transport2 is sse_transport assert mcp_server2 is mcp_server # Yield control to the event loop to satisfy async usage in this test await asyncio.sleep(0) async def test_init_mcp_servers(user_test_project, other_test_project): """Test the initialization of MCP servers for all projects.""" # Clear existing caches project_sse_transports.clear() project_mcp_servers.clear() # Test the initialization function await init_mcp_servers() # Verify that both test projects have SSE transports and MCP servers initialized project1_id = str(user_test_project.id) project2_id = str(other_test_project.id) # Both projects should have SSE transports created assert project1_id in project_sse_transports assert project2_id in project_sse_transports # Both projects should have MCP servers created assert project1_id in project_mcp_servers assert project2_id in project_mcp_servers # Verify the correct configuration assert isinstance(project_sse_transports[project1_id], SseServerTransport) assert isinstance(project_sse_transports[project2_id], SseServerTransport) assert project_mcp_servers[project1_id].project_id == user_test_project.id assert project_mcp_servers[project2_id].project_id == other_test_project.id async def test_init_mcp_servers_error_handling(): """Test that init_mcp_servers handles errors correctly and continues initialization.""" # Clear existing caches project_sse_transports.clear() project_mcp_servers.clear() # Create a mock to simulate an error when initializing one project original_get_project_sse = get_project_sse def mock_get_project_sse(project_id): # Raise an exception for the first project only if not project_sse_transports: # Only for the first project msg = "Test error for project SSE creation" raise ValueError(msg) return original_get_project_sse(project_id) # Apply the patch with patch("langflow.api.v1.mcp_projects.get_project_sse", side_effect=mock_get_project_sse): # This should not raise any exception, as the error should be caught await init_mcp_servers() @pytest.mark.asyncio async def test_mcp_longterm_token_fails_without_superuser(): """When AUTO_LOGIN is false and no superuser exists, creating a long-term token should raise 400. This simulates a clean DB with AUTO_LOGIN disabled and without provisioning a superuser. """ settings_service = get_settings_service() settings_service.auth_settings.AUTO_LOGIN = False # Ensure no superuser exists in DB async with get_db_service().with_session() as session: result = await session.exec(select(User).where(User.is_superuser == True)) # noqa: E712 users = result.all() for user in users: await session.delete(user) await session.commit() # Now attempt to create long-term token -> expect HTTPException 400 async with get_db_service().with_session() as session: with pytest.raises(HTTPException, match="Auto login required to create a long-term token"): await create_user_longterm_token(session)