"""Tests for OAuth proxy with persistent storage.""" import tempfile import warnings from collections.abc import AsyncGenerator from pathlib import Path from unittest.mock import AsyncMock, Mock import pytest from inline_snapshot import snapshot from key_value.aio.protocols import AsyncKeyValue from key_value.aio.stores.filetree import FileTreeStore from key_value.aio.stores.memory import MemoryStore from mcp.shared.auth import OAuthClientInformationFull from pydantic import AnyUrl from fastmcp.server.auth.auth import TokenVerifier from fastmcp.server.auth.oauth_proxy import OAuthProxy class TestOAuthProxyStorage: """Tests for OAuth proxy client storage functionality.""" @pytest.fixture def jwt_verifier(self): """Create a mock JWT verifier.""" verifier = Mock() verifier.required_scopes = ["read", "write"] verifier.verify_token = AsyncMock(return_value=None) return verifier @pytest.fixture async def temp_storage(self) -> AsyncGenerator[FileTreeStore, None]: """Create file-based storage for testing.""" with tempfile.TemporaryDirectory() as temp_dir: with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) yield FileTreeStore(data_directory=Path(temp_dir)) @pytest.fixture def memory_storage(self) -> MemoryStore: """Create in-memory storage for testing.""" return MemoryStore() def create_proxy( self, jwt_verifier: TokenVerifier, storage: AsyncKeyValue | None = None ) -> OAuthProxy: """Create an OAuth proxy with specified storage.""" return OAuthProxy( upstream_authorization_endpoint="https://github.com/login/oauth/authorize", upstream_token_endpoint="https://github.com/login/oauth/access_token", upstream_client_id="test-client-id", upstream_client_secret="test-client-secret", token_verifier=jwt_verifier, base_url="https://myserver.com", redirect_path="/auth/callback", client_storage=storage, jwt_signing_key="test-secret", ) async def test_register_and_get_client(self, jwt_verifier, temp_storage): """Test registering and retrieving a client.""" proxy = self.create_proxy(jwt_verifier, storage=temp_storage) # Register client client_info = OAuthClientInformationFull( client_id="test-client-123", client_secret="secret-456", redirect_uris=[AnyUrl("http://localhost:8080/callback")], grant_types=["authorization_code", "refresh_token"], scope="read write", ) await proxy.register_client(client_info) # Get client back client = await proxy.get_client("test-client-123") assert client is not None assert client.client_id == "test-client-123" # Proxy uses token_endpoint_auth_method="none", so client_secret is not stored assert client.client_secret is None assert client.scope == "read write" async def test_client_persists_across_proxy_instances( self, jwt_verifier: TokenVerifier, temp_storage: AsyncKeyValue ): """Test that clients persist when proxy is recreated.""" # First proxy registers client proxy1 = self.create_proxy(jwt_verifier, storage=temp_storage) client_info = OAuthClientInformationFull( client_id="persistent-client", client_secret="persistent-secret", redirect_uris=[AnyUrl("http://localhost:9999/callback")], scope="openid profile", ) await proxy1.register_client(client_info) # Second proxy can retrieve it proxy2 = self.create_proxy(jwt_verifier, storage=temp_storage) client = await proxy2.get_client("persistent-client") assert client is not None # Proxy uses token_endpoint_auth_method="none", so client_secret is not stored assert client.client_secret is None assert client.scope == "openid profile" async def test_nonexistent_client_returns_none( self, jwt_verifier: TokenVerifier, temp_storage: AsyncKeyValue ): """Test that requesting non-existent client returns None.""" proxy = self.create_proxy(jwt_verifier, storage=temp_storage) client = await proxy.get_client("does-not-exist") assert client is None async def test_proxy_dcr_client_redirect_validation( self, jwt_verifier: TokenVerifier, temp_storage: AsyncKeyValue ): """Test that OAuthProxyClient is created with redirect URI patterns.""" proxy = OAuthProxy( upstream_authorization_endpoint="https://github.com/login/oauth/authorize", upstream_token_endpoint="https://github.com/login/oauth/access_token", upstream_client_id="test-client-id", upstream_client_secret="test-client-secret", token_verifier=jwt_verifier, base_url="https://myserver.com", allowed_client_redirect_uris=["http://localhost:*"], client_storage=temp_storage, jwt_signing_key="test-secret", ) client_info = OAuthClientInformationFull( client_id="test-proxy-client", client_secret="secret", redirect_uris=[AnyUrl("http://localhost:8080/callback")], ) await proxy.register_client(client_info) # Get client back - should be OAuthProxyClient client = await proxy.get_client("test-proxy-client") assert client is not None # OAuthProxyClient should validate dynamic localhost ports validated = client.validate_redirect_uri( AnyUrl("http://localhost:12345/callback") ) assert validated is not None async def test_in_memory_storage_option(self, jwt_verifier): """Test using in-memory storage explicitly.""" storage = MemoryStore() proxy = self.create_proxy(jwt_verifier, storage=storage) client_info = OAuthClientInformationFull( client_id="memory-client", client_secret="memory-secret", redirect_uris=[AnyUrl("http://localhost:8080/callback")], ) await proxy.register_client(client_info) client = await proxy.get_client("memory-client") assert client is not None # Create new proxy with same storage instance proxy2 = self.create_proxy(jwt_verifier, storage=storage) client2 = await proxy2.get_client("memory-client") assert client2 is not None # But new storage instance won't have it proxy3 = self.create_proxy(jwt_verifier, storage=MemoryStore()) client3 = await proxy3.get_client("memory-client") assert client3 is None async def test_storage_data_structure(self, jwt_verifier, temp_storage): """Test that storage uses proper structured format.""" proxy = self.create_proxy(jwt_verifier, storage=temp_storage) client_info = OAuthClientInformationFull( client_id="structured-client", client_secret="secret", redirect_uris=[AnyUrl("http://localhost:8080/callback")], ) await proxy.register_client(client_info) # Check raw storage data raw_data = await temp_storage.get( collection="mcp-oauth-proxy-clients", key="structured-client" ) assert raw_data is not None assert raw_data == snapshot( { "redirect_uris": ["http://localhost:8080/callback"], "token_endpoint_auth_method": "none", "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], "scope": "read write", "client_name": None, "client_uri": None, "logo_uri": None, "contacts": None, "tos_uri": None, "policy_uri": None, "jwks_uri": None, "jwks": None, "software_id": None, "software_version": None, "application_type": "native", "issuer": None, "client_id": "structured-client", "client_secret": None, "client_id_issued_at": None, "client_secret_expires_at": None, "allowed_redirect_uri_patterns": None, "cimd_document": None, "cimd_fetched_at": None, } )