chore: import upstream snapshot with attribution

tests/unittests/tools/mcp_tool/test_mcp_toolset_auth.py

@@ -0,0 +1,298 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for MCPToolset authentication functionality."""
+
+import base64
+from unittest.mock import Mock
+
+from fastapi.openapi.models import APIKey as APIKeyScheme
+from fastapi.openapi.models import APIKeyIn
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import HttpAuth
+from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
+from mcp import StdioServerParameters
+import pytest
+
+
+class TestMcpToolsetGetAuthConfig:
+  """Tests for McpToolset.get_auth_config method."""
+
+  def test_get_auth_config_returns_none_without_auth_scheme(self):
+    """Test that get_auth_config returns None when no auth configured."""
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"])
+    )
+
+    assert toolset.get_auth_config() is None
+
+  def test_get_auth_config_returns_config_with_auth_scheme(self):
+    """Test that get_auth_config returns AuthConfig when auth configured."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access"},
+            )
+        )
+    )
+    auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+    auth_config = toolset.get_auth_config()
+    assert auth_config is not None
+    assert auth_config.auth_scheme == auth_scheme
+    assert auth_config.raw_auth_credential == auth_credential
+
+  def test_get_auth_config_returns_same_instance(self):
+    """Test that get_auth_config returns the same instance each time."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Should return the same instance
+    config1 = toolset.get_auth_config()
+    config2 = toolset.get_auth_config()
+    assert config1 is config2
+
+
+class TestMcpToolsetGetAuthHeaders:
+  """Tests for McpToolset._get_auth_headers method."""
+
+  @pytest.fixture
+  def toolset_with_oauth2(self):
+    """Create a toolset with OAuth2 auth configured."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access"},
+            )
+        )
+    )
+    auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    return McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+    )
+
+  def test_get_auth_headers_returns_none_without_auth_config(self):
+    """Test that _get_auth_headers returns None without auth config."""
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"])
+    )
+
+    assert toolset._get_auth_headers() is None
+
+  def test_get_auth_headers_returns_none_without_exchanged_credential(
+      self, toolset_with_oauth2
+  ):
+    """Test that _get_auth_headers returns None without exchanged credential."""
+    # No exchanged credential set yet
+    assert toolset_with_oauth2._get_auth_headers() is None
+
+  def test_get_auth_headers_oauth2_bearer_token(self, toolset_with_oauth2):
+    """Test that _get_auth_headers returns Bearer token for OAuth2."""
+    # Set exchanged credential with access token
+    toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="test-access-token"),
+    )
+
+    headers = toolset_with_oauth2._get_auth_headers()
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer test-access-token"
+
+  def test_get_auth_headers_http_bearer_token(self):
+    """Test that _get_auth_headers returns Bearer token for HTTP bearer."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with HTTP bearer token
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="bearer",
+            credentials=HttpCredentials(token="test-bearer-token"),
+        ),
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer test-bearer-token"
+
+  def test_get_auth_headers_http_basic_auth(self):
+    """Test that _get_auth_headers returns Basic auth for HTTP basic."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={},
+            )
+        )
+    )
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with HTTP basic auth
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        http=HttpAuth(
+            scheme="basic",
+            credentials=HttpCredentials(
+                username="testuser",
+                password="testpass",
+            ),
+        ),
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    expected_credentials = base64.b64encode(b"testuser:testpass").decode()
+    assert headers["Authorization"] == f"Basic {expected_credentials}"
+
+  def test_get_auth_headers_api_key_header(self):
+    """Test that _get_auth_headers returns API key in header."""
+    # Note: fastapi's APIKey model uses 'in' not 'in_', but accepts both
+    auth_scheme = APIKeyScheme(**{
+        "in": APIKeyIn.header,
+        "name": "X-API-Key",
+    })
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with API key
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY,
+        api_key="test-api-key-12345",
+    )
+
+    headers = toolset._get_auth_headers()
+
+    assert headers is not None
+    assert headers["X-API-Key"] == "test-api-key-12345"
+
+  def test_get_auth_headers_api_key_non_header_logs_warning(self, caplog):
+    """Test that non-header API key logs a warning."""
+    # Note: fastapi's APIKey model uses 'in' not 'in_'
+    auth_scheme = APIKeyScheme(**{
+        "in": APIKeyIn.query,  # Query param, not header
+        "name": "api_key",
+    })
+
+    toolset = McpToolset(
+        connection_params=StdioServerParameters(command="echo", args=["test"]),
+        auth_scheme=auth_scheme,
+    )
+
+    # Set exchanged credential with API key
+    toolset._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY,
+        api_key="test-api-key",
+    )
+
+    headers = toolset._get_auth_headers()
+
+    # Should return None for non-header API key
+    assert headers is None
+
+  def test_get_auth_headers_reads_from_readonly_context(
+      self, toolset_with_oauth2
+  ):
+    """Test that _get_auth_headers reads from ReadonlyContext first."""
+    from google.adk.agents.readonly_context import ReadonlyContext
+
+    mock_readonly_context = Mock(spec=ReadonlyContext)
+    mock_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-context"),
+    )
+    mock_readonly_context.get_credential.return_value = mock_credential
+
+    # Even if exchanged_auth_credential has a different value
+    toolset_with_oauth2._auth_config.exchanged_auth_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="token-from-config"),
+    )
+
+    headers = toolset_with_oauth2._get_auth_headers(
+        readonly_context=mock_readonly_context
+    )
+
+    assert headers is not None
+    assert headers["Authorization"] == "Bearer token-from-context"
+    mock_readonly_context.get_credential.assert_called_once_with(
+        toolset_with_oauth2._auth_config.credential_key
+    )