From dc0fc95519233182a5bfa58d1081f635d0de7837 Mon Sep 17 00:00:00 2001 From: Kamil Potrec Date: Wed, 1 Oct 2025 22:26:24 +0100 Subject: [PATCH] feat(keycardai-mcp-fastmcp): ability to mock internal access context for testing --- packages/mcp-fastmcp/README.md | 161 ++++++++++++++++++ .../mcp/integrations/fastmcp/__init__.py | 5 +- .../mcp/integrations/fastmcp/provider.py | 100 +++++------ .../integrations/fastmcp/testing/__init__.py | 36 ++++ .../fastmcp/testing/test_utils.py | 98 +++++++++++ 5 files changed, 349 insertions(+), 51 deletions(-) create mode 100644 packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/__init__.py create mode 100644 packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/test_utils.py diff --git a/packages/mcp-fastmcp/README.md b/packages/mcp-fastmcp/README.md index 963c898e..bb50af55 100644 --- a/packages/mcp-fastmcp/README.md +++ b/packages/mcp-fastmcp/README.md @@ -267,6 +267,167 @@ auth_provider = AuthProvider( ) ``` +## Testing + +This section provides comprehensive guidance on testing your FastMCP servers that use Keycard authentication. The examples show how to use the `mock_access_context` utility to easily mock authentication without needing to understand the internal SDK implementation. + +### Overview + +When testing FastMCP servers with Keycard authentication, you need to mock the authentication system. The `mock_access_context` utility provides four main testing scenarios: + +1. **Default token** - Always returns a default access token for any resource +2. **Custom token** - Returns a specific access token for any resource +3. **Resource-specific tokens** - Returns different tokens for different resources +4. **Error scenarios** - Simulates authentication failures + +### Basic Test Setup + +### Testing Tools With Grant Decorators + +For tools that use the `@grant` decorator, use the `mock_access_context` utility to mock the authentication system: + +#### 1. Default Token (Simple Case) + +```python +@pytest.mark.asyncio +async def test_tool_with_default_token(auth_provider): + """Test a tool with default access token.""" + + # Create FastMCP server + mcp = FastMCP("Test Server", auth=auth_provider.get_remote_auth_provider()) + + @mcp.tool() + @auth_provider.grant("https://api.example.com") + def call_external_api(ctx: Context, query: str) -> str: + access_context = ctx.get_state("keycardai") + + if access_context.has_errors(): + return f"Error: {access_context.get_errors()}" + + token = access_context.access("https://api.example.com").access_token + return f"API result for {query} with token {token}" + + # Test with default token + with mock_access_context(): # Uses "test_access_token" by default + async with Client(mcp) as client: + result = await client.call_tool("call_external_api", {"query": "test"}) + + assert result is not None + assert "test_access_token" in result.data + assert "API result for test" in result.data +``` + +#### 2. Custom Token + +```python +@pytest.mark.asyncio +async def test_tool_with_custom_token(auth_provider): + """Test a tool with a specific access token.""" + + mcp = FastMCP("Test Server", auth=auth_provider.get_remote_auth_provider()) + + @mcp.tool() + @auth_provider.grant("https://api.example.com") + def call_external_api(ctx: Context, query: str) -> str: + access_context = ctx.get_state("keycardai") + token = access_context.access("https://api.example.com").access_token + return f"API result for {query} with token {token}" + + # Test with custom token + with mock_access_context(access_token="my_custom_token_123"): + async with Client(mcp) as client: + result = await client.call_tool("call_external_api", {"query": "test"}) + + assert "my_custom_token_123" in result.data +``` + +#### 3. Resource-Specific Tokens + +```python +@pytest.mark.asyncio +async def test_tool_with_resource_specific_tokens(auth_provider): + """Test a tool with different tokens for different resources.""" + + mcp = FastMCP("Test Server", auth=auth_provider.get_remote_auth_provider()) + + @mcp.tool() + @auth_provider.grant(["https://api.example.com", "https://calendar-api.com"]) + def sync_data(ctx: Context) -> str: + access_context = ctx.get_state("keycardai") + + api_token = access_context.access("https://api.example.com").access_token + calendar_token = access_context.access("https://calendar-api.com").access_token + + return f"API: {api_token}, Calendar: {calendar_token}" + + # Test with resource-specific tokens + with mock_access_context(resource_tokens={ + "https://api.example.com": "api_token_123", + "https://calendar-api.com": "calendar_token_456" + }): + async with Client(mcp) as client: + result = await client.call_tool("sync_data", {}) + + assert "api_token_123" in result.data + assert "calendar_token_456" in result.data +``` + +### Testing Error Scenarios + +Test how your tools handle authentication errors using the `has_errors` parameter: + +```python +@pytest.mark.asyncio +async def test_tool_with_authentication_error(auth_provider): + """Test tool behavior when authentication fails.""" + + mcp = FastMCP("Test Server", auth=auth_provider.get_remote_auth_provider()) + + @mcp.tool() + @auth_provider.grant("https://api.example.com") + def failing_tool(ctx: Context, query: str) -> str: + access_context = ctx.get_state("keycardai") + + # Always check for errors first + if access_context.has_errors(): + return f"Authentication failed: {access_context.get_errors()}" + + token = access_context.access("https://api.example.com").access_token + return f"Success: {query}" + + # Test with authentication error + with mock_access_context(has_errors=True, error_message="Token exchange failed"): + async with Client(mcp) as client: + result = await client.call_tool("failing_tool", {"query": "test"}) + + assert result is not None + assert "Authentication failed" in result.data + assert "Token exchange failed" in result.data + +@pytest.mark.asyncio +async def test_tool_with_custom_error_message(auth_provider): + """Test tool with custom error message.""" + + mcp = FastMCP("Test Server", auth=auth_provider.get_remote_auth_provider()) + + @mcp.tool() + @auth_provider.grant("https://api.example.com") + def error_handling_tool(ctx: Context) -> str: + access_context = ctx.get_state("keycardai") + + if access_context.has_errors(): + return f"Error occurred: {access_context.get_errors()}" + + return "Success" + + # Test with custom error message + with mock_access_context(has_errors=True, error_message="Custom auth error"): + async with Client(mcp) as client: + result = await client.call_tool("error_handling_tool", {}) + + assert "Custom auth error" in result.data +``` + ## Examples For complete examples and advanced usage patterns, see our [documentation](https://docs.keycard.ai). diff --git a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/__init__.py b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/__init__.py index 0dcd0086..ea4261e3 100644 --- a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/__init__.py +++ b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/__init__.py @@ -66,7 +66,6 @@ async def sync_calendar_to_drive(ctx: Context): ) """ -# Re-export commonly used auth strategies for convenience from keycardai.mcp.server.auth.client_factory import ClientFactory, DefaultClientFactory from keycardai.mcp.server.exceptions import ( # Specific exceptions @@ -88,6 +87,7 @@ async def sync_calendar_to_drive(ctx: Context): ) from .provider import AccessContext, AuthProvider +from .testing import mock_access_context __all__ = [ # Core classes @@ -115,4 +115,7 @@ async def sync_calendar_to_drive(ctx: Context): "ResourceAccessError", "TokenExchangeError", "MetadataDiscoveryError", + + # Testing utilities + "mock_access_context", ] diff --git a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/provider.py b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/provider.py index a786c508..a0fc25e5 100644 --- a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/provider.py +++ b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/provider.py @@ -404,81 +404,81 @@ async def my_async_tool(ctx: Context, user_id: str): - Preserves original function signature and behavior - Provides detailed error messages for debugging """ + def _has_context(func: Callable) -> bool: + sig = inspect.signature(func) + for value in sig.parameters.values(): + if value.annotation == Context: + return True + return False + + def _get_context(*args, **kwargs) -> Context | None: + for value in args: + if isinstance(value, Context): + return value + for value in kwargs.values(): + if isinstance(value, Context): + return value + return None + + def _set_error(error: dict[str, str], resource: str | None, access_context: AccessContext, ctx: Context): + """Helper to set error context and call function.""" + if resource: + access_context.set_resource_error(resource, error) + else: + access_context.set_error(error) + ctx.set_state("keycardai", access_context) + + async def _call_func(is_async_func: bool, func: Callable, *args, **kwargs): + if is_async_func: + return await func(*args, **kwargs) + else: + return func(*args, **kwargs) + def decorator(func: Callable) -> Callable: is_async_func = inspect.iscoroutinefunction(func) - def _has_context(func: Callable) -> bool: - sig = inspect.signature(func) - for value in sig.parameters.values(): - if value.annotation == Context: - return True - return False - if not _has_context(func): raise MissingContextError() - def _get_context(*args, **kwargs) -> Context | None: - for value in args: - if isinstance(value, Context): - return value - for value in kwargs.values(): - if isinstance(value, Context): - return value - return None - - def _set_error(error: dict[str, str], resource: str | None, access_context: AccessContext, ctx: Context): - """Helper to set error context and call function.""" - if resource: - access_context.set_resource_error(resource, error) - else: - access_context.set_error(error) - ctx.set_state("keycardai", access_context) - - async def _call_func(func: Callable, *args, **kwargs): - if is_async_func: - return await func(*args, **kwargs) - else: - return func(*args, **kwargs) - @wraps(func) async def wrapper(*args, **kwargs) -> Any: - ctx = _get_context(*args, **kwargs) - if ctx is None: + _ctx = _get_context(*args, **kwargs) + if _ctx is None: raise MissingContextError() - access_context = AccessContext() + _access_context = AccessContext() try: - user_token = get_access_token() - if not user_token: + _user_token = get_access_token() + if not _user_token: _set_error({ "error": "No authentication token available. Please ensure you're properly authenticated.", - }, None, access_context, ctx) - return await _call_func(func, *args, **kwargs) + }, None, _access_context, _ctx) + return await _call_func(is_async_func, func, *args, **kwargs) except Exception as e: _set_error({ "error": "Failed to get access token from context. Ensure the Context parameter is properly annotated.", "raw_error": str(e), - }, None, access_context, ctx) - return await _call_func(func, *args, **kwargs) - resource_list = [resources] if isinstance(resources, str) else resources - access_tokens = {} - for resource in resource_list: + }, None, _access_context, _ctx) + return await _call_func(is_async_func, func, *args, **kwargs) + _resource_list = [resources] if isinstance(resources, str) else resources + _access_tokens = {} + for resource in _resource_list: try: - token_response = await self.client.exchange_token( - subject_token=user_token.token, + _token_response = await self.client.exchange_token( + subject_token=_user_token.token, resource=resource, subject_token_type="urn:ietf:params:oauth:token-type:access_token" ) - access_tokens[resource] = token_response + _access_tokens[resource] = _token_response except Exception as e: _set_error({ "error": f"Token exchange failed for {resource}: {e}", "raw_error": str(e), - }, resource, access_context, ctx) - return await _call_func(func, *args, **kwargs) + }, resource, _access_context, _ctx) + return await _call_func(is_async_func, func, *args, **kwargs) # Set successful tokens on the existing access_context (preserves any resource errors) - access_context.set_bulk_tokens(access_tokens) - ctx.set_state("keycardai", access_context) - return await _call_func(func, *args, **kwargs) + _access_context.set_bulk_tokens(_access_tokens) + _ctx.set_state("keycardai", _access_context) + return await _call_func(is_async_func, func, *args, **kwargs) return wrapper return decorator diff --git a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/__init__.py b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/__init__.py new file mode 100644 index 00000000..06af6099 --- /dev/null +++ b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/__init__.py @@ -0,0 +1,36 @@ +"""Testing utilities for FastMCP integration with Keycard authentication. + +This module provides mock implementations and utilities for testing FastMCP servers +that use Keycard authentication without requiring real OAuth flows or network calls. + +Components: +- mock_access_context: Context manager for mocking authentication in tests + +Example: + from keycardai.mcp.integrations.fastmcp.testing import mock_access_context + + # Test successful authentication with default token + with mock_access_context(): + # Your test code here - will return "test_access_token" for any resource + + # Test with specific access token + with mock_access_context(access_token="my_custom_token"): + # Your test code here - will return "my_custom_token" for any resource + + # Test with resource-specific tokens + with mock_access_context(resource_tokens={ + "https://api.example.com": "token_123", + "https://api.other.com": "token_456" + }): + # Your test code here - will return specific tokens for each resource + + # Test error scenarios + with mock_access_context(has_errors=True, error_message="Auth failed"): + # Your test code here - access_context.has_errors() will return True +""" + +from .test_utils import mock_access_context + +__all__ = [ + "mock_access_context", +] diff --git a/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/test_utils.py b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/test_utils.py new file mode 100644 index 00000000..2c241d70 --- /dev/null +++ b/packages/mcp-fastmcp/src/keycardai/mcp/integrations/fastmcp/testing/test_utils.py @@ -0,0 +1,98 @@ +from contextlib import contextmanager +from unittest.mock import Mock, patch + +from keycardai.oauth.types.models import TokenResponse + + +@contextmanager +def mock_access_context( + access_token: str = "test_access_token", + resource_tokens: dict[str, str] | None = None, + has_errors: bool = False, + error_message: str = "Mock authentication error", +): + """Mock the authentication system for testing. + + Args: + access_token: Default access token to return for any resource (str) + resource_tokens: Dict mapping resource URLs to specific access tokens (dict[str, str]) + has_errors: Whether the access context should report errors (bool) + error_message: Error message to return when has_errors=True (str) + + Examples: + # 1. Default - always returns access token + with mock_access_context(): + # Will return "test_access_token" for any resource + + # 2. Returns access token for provided resource + with mock_access_context(access_token="my_token"): + # Will return "my_token" for any resource + + # 3. Return access token for provided dict of resources + with mock_access_context(resource_tokens={ + "https://api.example.com": "token_123", + "https://api.other.com": "token_456" + }): + # Will return specific tokens for each resource + # Any resource not in the dict will set has_errors=True with "Resource not granted" message + + # 4. Returns error set to true and error message + with mock_access_context(has_errors=True, error_message="Auth failed"): + # Will report errors with the specified message + """ + with patch('keycardai.mcp.integrations.fastmcp.provider.AccessContext') as mock_access_context_class, \ + patch('keycardai.mcp.integrations.fastmcp.provider.get_access_token') as mock_get_access_token: + + mock_access_context_instance = Mock() + mock_access_context_instance.has_errors.return_value = has_errors + + if has_errors: + # Return proper error structure matching AccessContext.get_errors() + mock_access_context_instance.get_errors.return_value = { + "resource_errors": {}, + "error": {"error": error_message} + } + mock_access_context_instance.access.side_effect = Exception(error_message) + else: + def mock_access_method(resource_url): + if resource_tokens is not None: + if resource_url in resource_tokens: + # Return proper TokenResponse object + return TokenResponse( + access_token=resource_tokens[resource_url], + token_type="Bearer" + ) + else: + # Resource not granted - set error state and raise exception + mock_access_context_instance.has_errors.return_value = True + mock_access_context_instance.get_errors.return_value = { + "resource_errors": { + resource_url: {"error": f"Resource not granted: {resource_url}"} + }, + "error": None + } + from keycardai.mcp.server.exceptions import ResourceAccessError + raise ResourceAccessError() + else: + # Return proper TokenResponse object + return TokenResponse( + access_token=access_token, + token_type="Bearer" + ) + + mock_access_context_instance.access = mock_access_method + mock_access_context_instance.get_errors.return_value = { + "resource_errors": {}, + "error": None + } + + mock_access_context_instance.set_bulk_tokens = Mock() + mock_access_context_instance.set_error = Mock() + mock_access_context_instance.set_resource_error = Mock() + mock_access_context_class.return_value = mock_access_context_instance + + mock_user_token = Mock() + mock_user_token.token = "user_jwt_token" + mock_get_access_token.return_value = mock_user_token + + yield mock_access_context_instance