Added unit-tests

Author: yurikunash

src/mcp/client/auth.py

@@ -34,31 +34,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _extract_resource_metadata_from_www_auth(response: httpx.Response) -> str | None:
-    """
-    Extract protected resource metadata URL from WWW-Authenticate header as per RFC9728.
-        
-    Returns:
-        Resource metadata URL if found in WWW-Authenticate header, None otherwise
-    """
-    if not response or response.status_code != 401:
-        return None
-        
-    www_auth_header = response.headers.get("WWW-Authenticate")
-    if not www_auth_header:
-        return None
-    
-    # Pattern matches: resource_metadata="url" or resource_metadata=url (unquoted)
-    pattern = r'resource_metadata=(?:"([^"]+)"|([^\s,]+))'
-    match = re.search(pattern, www_auth_header)
-    
-    if match:
-        # Return quoted value if present, otherwise unquoted value
-        return match.group(1) or match.group(2)
-    
-    return None
-
-
 class OAuthFlowError(Exception):
     """Base exception for OAuth flow errors."""
 
@@ -229,9 +204,33 @@ def __init__(
         )
         self._initialized = False
 
-    async def _discover_protected_resource(self, response: httpx.Response | None = None) -> httpx.Request:
-        # RFC9728: Try to extract resource_metadata URL from WWW-Authenticate header
-        url = _extract_resource_metadata_from_www_auth(response) if response else None
+    def _extract_resource_metadata_from_www_auth(self, init_response: httpx.Response) -> str | None:
+        """
+        Extract protected resource metadata URL from WWW-Authenticate header as per RFC9728.
+            
+        Returns:
+            Resource metadata URL if found in WWW-Authenticate header, None otherwise
+        """
+        if not init_response or init_response.status_code != 401:
+            return None
+            
+        www_auth_header = init_response.headers.get("WWW-Authenticate")
+        if not www_auth_header:
+            return None
+        
+        # Pattern matches: resource_metadata="url" or resource_metadata=url (unquoted)
+        pattern = r'resource_metadata=(?:"([^"]+)"|([^\s,]+))'
+        match = re.search(pattern, www_auth_header)
+        
+        if match:
+            # Return quoted value if present, otherwise unquoted value
+            return match.group(1) or match.group(2)
+        
+        return None
+
+    async def _discover_protected_resource(self, init_response: httpx.Response | None = None) -> httpx.Request:
+        # RFC9728: Try to extract resource_metadata URL from WWW-Authenticate header of the initial response
+        url = self._extract_resource_metadata_from_www_auth(init_response) if init_response else None
 
         if not url:
             # Fallback to well-known discovery

tests/client/test_auth.py

@@ -196,13 +196,47 @@ class TestOAuthFlow:
     """Test OAuth flow methods."""
 
     @pytest.mark.anyio
-    async def test_discover_protected_resource_request(self, oauth_provider):
-        """Test protected resource discovery request building."""
-        request = await oauth_provider._discover_protected_resource()
+    async def test_discover_protected_resource_request(self, client_metadata, mock_storage):
+        """Test protected resource discovery request building maintains backward compatibility."""
+        async def redirect_handler(url: str) -> None:
+            pass
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", "test_state"
 
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
+        )
+        
+        # Test without response (backward compatibility)
+        request = await provider._discover_protected_resource()
         assert request.method == "GET"
         assert str(request.url) == "https://api.example.com/.well-known/oauth-protected-resource"
         assert "mcp-protocol-version" in request.headers
+        
+        # Test with response but no WWW-Authenticate (fallback)
+        init_response = httpx.Response(
+            status_code=401,
+            headers={},
+            request=httpx.Request("GET", "https://request-api.example.com")
+        )
+
+        request = await provider._discover_protected_resource(init_response)
+        assert request.method == "GET" 
+        assert str(request.url) == "https://api.example.com/.well-known/oauth-protected-resource"
+        assert "mcp-protocol-version" in request.headers
+
+        # Test with WWW-Authenticate header
+        init_response.headers["WWW-Authenticate"] = 'Bearer resource_metadata="https://prm.example.com/.well-known/oauth-protected-resource/path"'
+        
+        request = await provider._discover_protected_resource(init_response)
+        assert request.method == "GET" 
+        assert str(request.url) == "https://prm.example.com/.well-known/oauth-protected-resource/path"
+        assert "mcp-protocol-version" in request.headers
 
     @pytest.mark.anyio
     async def test_discover_oauth_metadata_request(self, oauth_provider):
@@ -544,3 +578,108 @@ async def test_auth_flow_with_valid_tokens(self, oauth_provider, mock_storage, v
             await auth_flow.asend(response)
         except StopAsyncIteration:
             pass  # Expected
+
+
+class TestRFC9728WWWAuthenticate:
+    """Test RFC9728 WWW-Authenticate header parsing functionality."""
+
+    @pytest.mark.parametrize("www_auth_header,expected_url", [
+        # Quoted URL
+        ('Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"', 
+         "https://api.example.com/.well-known/oauth-protected-resource"),
+        # Unquoted URL
+        ("Bearer resource_metadata=https://api.example.com/.well-known/oauth-protected-resource", 
+         "https://api.example.com/.well-known/oauth-protected-resource"),
+        # Complex header with multiple parameters
+        ('Bearer realm="api", resource_metadata="https://api.example.com/.well-known/oauth-protected-resource", error="insufficient_scope"', 
+         "https://api.example.com/.well-known/oauth-protected-resource"),
+        # Different URL format
+        ('Bearer resource_metadata="https://custom.domain.com/metadata"', 
+         "https://custom.domain.com/metadata"),
+        # With path and query params
+        ('Bearer resource_metadata="https://api.example.com/auth/metadata?version=1"', 
+         "https://api.example.com/auth/metadata?version=1"),
+    ])
+    def test_extract_resource_metadata_from_www_auth_valid_cases(self, client_metadata, mock_storage, www_auth_header, expected_url):
+        """Test extraction of resource_metadata URL from various valid WWW-Authenticate headers."""
+        async def redirect_handler(url: str) -> None:
+            pass
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", "test_state"
+
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com/v1/mcp",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
+        )
+        
+        init_response = httpx.Response(
+            status_code=401,
+            headers={"WWW-Authenticate": www_auth_header},
+            request=httpx.Request("GET", "https://api.example.com/test")
+        )
+        
+        result = provider._extract_resource_metadata_from_www_auth(init_response)
+        assert result == expected_url
+
+    @pytest.mark.parametrize("status_code,www_auth_header,description", [
+        # No header
+        (401, None, "no WWW-Authenticate header"),
+        # Empty header  
+        (401, "", "empty WWW-Authenticate header"),
+        # Header without resource_metadata
+        (401, 'Bearer realm="api", error="insufficient_scope"', "no resource_metadata parameter"),
+        # Malformed header
+        (401, "Bearer resource_metadata=", "malformed resource_metadata parameter"),
+        # Non-401 status code
+        (200, 'Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"', "200 OK response"),
+        (500, 'Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"', "500 error response"),
+    ])
+    def test_extract_resource_metadata_from_www_auth_invalid_cases(self, client_metadata, mock_storage, status_code, www_auth_header, description):
+        """Test extraction returns None for invalid cases."""
+        async def redirect_handler(url: str) -> None:
+            pass
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", "test_state"
+
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com/v1/mcp",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
+        )
+        
+        headers = {"WWW-Authenticate": www_auth_header} if www_auth_header is not None else {}
+        init_response = httpx.Response(
+            status_code=status_code,
+            headers=headers,
+            request=httpx.Request("GET", "https://api.example.com/test")
+        )
+        
+        result = provider._extract_resource_metadata_from_www_auth(init_response)
+        assert result is None, f"Should return None for {description}"
+
+    def test_extract_resource_metadata_from_www_auth_none_response(self, client_metadata, mock_storage):
+        """Test extraction with None response returns None."""
+        async def redirect_handler(url: str) -> None:
+            pass
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", "test_state"
+
+        provider = OAuthClientProvider(
+            server_url="https://api.example.com/v1/mcp",
+            client_metadata=client_metadata,
+            storage=mock_storage,
+            redirect_handler=redirect_handler,
+            callback_handler=callback_handler,
+        )
+        
+        result = provider._extract_resource_metadata_from_www_auth(None)
+        assert result is None
+