# 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.

import asyncio
import base64
from io import StringIO
import pickle
import sys
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import Mock

from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlows
from google.adk.agents.readonly_context import ReadonlyContext
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.load_mcp_resource_tool import LoadMcpResourceTool
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
from google.adk.tools.mcp_tool.mcp_tool import MCPTool
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
from google.adk.tools.tool_configs import ToolArgsConfig
from mcp import StdioServerParameters
from mcp.types import BlobResourceContents
from mcp.types import ListResourcesResult
from mcp.types import ReadResourceResult
from mcp.types import Resource
from mcp.types import TextResourceContents
import pytest


class MockMCPTool:
  """Mock MCP Tool for testing."""

  def __init__(self, name, description="Test tool description"):
    self.name = name
    self.description = description
    self.inputSchema = {
        "type": "object",
        "properties": {"param": {"type": "string"}},
    }


class MockListToolsResult:
  """Mock ListToolsResult for testing."""

  def __init__(self, tools):
    self.tools = tools


class TestMcpToolset:
  """Test suite for McpToolset class."""

  def setup_method(self):
    """Set up test fixtures."""
    self.mock_stdio_params = StdioServerParameters(
        command="test_command", args=[]
    )
    self.mock_session_manager = Mock(spec=MCPSessionManager)
    self.mock_session = AsyncMock()
    self.mock_session_manager.create_session = AsyncMock(
        return_value=self.mock_session
    )

  def test_init_basic(self):
    """Test basic initialization with StdioServerParameters."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)

    # Note: StdioServerParameters gets converted to StdioConnectionParams internally
    assert toolset._errlog == sys.stderr
    assert toolset._auth_scheme is None
    assert toolset._auth_credential is None
    assert toolset._use_mcp_resources is False

  def test_init_with_use_mcp_resources(self):
    """Test initialization with use_mcp_resources."""
    toolset = McpToolset(
        connection_params=self.mock_stdio_params, use_mcp_resources=True
    )
    assert toolset._use_mcp_resources is True

  def test_connection_params(self):
    """Test getting connection params."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    assert toolset.connection_params == self.mock_stdio_params

  def test_auth_scheme(self):
    """Test getting auth scheme."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    assert toolset.auth_scheme is None

  def test_auth_credential(self):
    """Test getting auth credential."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    assert toolset.auth_credential is None

  def test_error_log(self):
    """Test getting error log."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    assert toolset.errlog == sys.stderr

  def test_auth_scheme_with_value(self):
    """Test getting auth scheme when provided at initialization."""
    auth_scheme = OAuth2(
        flows=OAuthFlows(
            authorizationCode=OAuthFlowAuthorizationCode(
                authorizationUrl="https://example.com/auth",
                tokenUrl="https://example.com/token",
                scopes={"read": "Read access"},
            )
        )
    )
    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        auth_scheme=auth_scheme,
    )
    assert toolset.auth_scheme == auth_scheme

  def test_require_confirmation(self):
    """Test getting require_confirmation flag."""
    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        require_confirmation=True,
    )
    assert toolset.require_confirmation is True

  def test_header_provider(self):
    """Test getting header_provider."""
    mock_header_provider = Mock()
    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        header_provider=mock_header_provider,
    )
    assert toolset.header_provider == mock_header_provider

  def test_auth_credential_with_value(self):
    """Test getting auth credential when provided at initialization."""
    mock_credential = Mock(spec=AuthCredential)
    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        auth_credential=mock_credential,
    )
    assert toolset.auth_credential == mock_credential

  def test_init_with_stdio_connection_params(self):
    """Test initialization with StdioConnectionParams."""
    stdio_params = StdioConnectionParams(
        server_params=self.mock_stdio_params, timeout=10.0
    )
    toolset = McpToolset(connection_params=stdio_params)

    assert toolset._connection_params == stdio_params

  def test_init_with_sse_connection_params(self):
    """Test initialization with SseConnectionParams."""
    sse_params = SseConnectionParams(
        url="https://example.com/mcp", headers={"Authorization": "Bearer token"}
    )
    toolset = McpToolset(connection_params=sse_params)

    assert toolset._connection_params == sse_params

  def test_init_with_streamable_http_params(self):
    """Test initialization with StreamableHTTPConnectionParams."""
    http_params = StreamableHTTPConnectionParams(
        url="https://example.com/mcp",
        headers={"Content-Type": "application/json"},
    )
    toolset = McpToolset(connection_params=http_params)

    assert toolset._connection_params == http_params

  def test_init_with_tool_filter_list(self):
    """Test initialization with tool filter as list."""
    tool_filter = ["tool1", "tool2"]
    toolset = McpToolset(
        connection_params=self.mock_stdio_params, tool_filter=tool_filter
    )

    # The tool filter is stored in the parent BaseToolset class
    # We can verify it by checking the filtering behavior in get_tools
    assert toolset._is_tool_selected is not None

  def test_init_with_auth(self):
    """Test initialization with authentication."""
    # Create real auth scheme instances

    auth_scheme = OAuth2(flows={})

    auth_credential = AuthCredential(
        auth_type="oauth2",
        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
    )

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        auth_scheme=auth_scheme,
        auth_credential=auth_credential,
    )

    assert toolset._auth_scheme == auth_scheme
    assert toolset._auth_credential == auth_credential

  def test_init_with_auth_and_credential_key(self):
    """Test initialization with authentication and a custom credential_key."""

    auth_scheme = OAuth2(flows={})

    auth_credential = AuthCredential(
        auth_type="oauth2",
        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
    )

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        auth_scheme=auth_scheme,
        auth_credential=auth_credential,
        credential_key="my_custom_key",
    )

    assert toolset._auth_scheme == auth_scheme
    assert toolset._auth_credential == auth_credential
    assert toolset._auth_config.credential_key == "my_custom_key"

  def test_from_config_with_credential_key(self):
    """Test that from_config correctly parses credential_key."""

    auth_scheme = OAuth2(flows={})

    config = ToolArgsConfig(
        stdio_server_params=self.mock_stdio_params,
        auth_scheme=auth_scheme,
        credential_key="my_custom_key",
    )
    toolset = McpToolset.from_config(config, "")

    assert isinstance(toolset._auth_scheme, OAuth2)
    assert toolset._auth_config.credential_key == "my_custom_key"

  def test_init_missing_connection_params(self):
    """Test initialization with missing connection params raises error."""
    with pytest.raises(ValueError, match="Missing connection params"):
      McpToolset(connection_params=None)

  @pytest.mark.asyncio
  async def test_get_tools_basic(self):
    """Test getting tools without filtering."""
    # Mock tools from MCP server
    mock_tools = [
        MockMCPTool("tool1"),
        MockMCPTool("tool2"),
        MockMCPTool("tool3"),
    ]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )

    toolset = McpToolset(
        connection_params=self.mock_stdio_params, use_mcp_resources=True
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools()

    assert len(tools) == 4
    for tool in tools[:3]:
      assert isinstance(tool, MCPTool)
    assert isinstance(tools[3], LoadMcpResourceTool)
    assert tools[0].name == "tool1"
    assert tools[1].name == "tool2"
    assert tools[2].name == "tool3"
    assert tools[3].name == "load_mcp_resource"

  @pytest.mark.asyncio
  async def test_get_tools_with_list_filter(self):
    """Test getting tools with list-based filtering."""
    # Mock tools from MCP server
    mock_tools = [
        MockMCPTool("tool1"),
        MockMCPTool("tool2"),
        MockMCPTool("tool3"),
    ]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )

    tool_filter = ["tool1", "tool3"]
    toolset = McpToolset(
        connection_params=self.mock_stdio_params, tool_filter=tool_filter
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools()

    assert len(tools) == 2
    assert tools[0].name == "tool1"
    assert tools[1].name == "tool3"

  @pytest.mark.asyncio
  async def test_get_tools_with_function_filter(self):
    """Test getting tools with function-based filtering."""
    # Mock tools from MCP server
    mock_tools = [
        MockMCPTool("read_file"),
        MockMCPTool("write_file"),
        MockMCPTool("list_directory"),
    ]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )

    def file_tools_filter(tool, context):
      """Filter for file-related tools only."""
      return "file" in tool.name

    toolset = McpToolset(
        connection_params=self.mock_stdio_params, tool_filter=file_tools_filter
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools()

    assert len(tools) == 2
    assert tools[0].name == "read_file"
    assert tools[1].name == "write_file"

  @pytest.mark.asyncio
  async def test_get_tools_with_header_provider(self):
    """Test get_tools with a header_provider."""
    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )
    mock_readonly_context = Mock(spec=ReadonlyContext)
    expected_headers = {"X-Tenant-ID": "test-tenant"}
    header_provider = Mock(return_value=expected_headers)

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        header_provider=header_provider,
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools(readonly_context=mock_readonly_context)

    assert len(tools) == 2
    header_provider.assert_called_once_with(mock_readonly_context)
    self.mock_session_manager.create_session.assert_called_once_with(
        headers=expected_headers
    )

  @pytest.mark.asyncio
  async def test_get_tools_with_async_header_provider(self):
    """Test get_tools with an async header_provider."""
    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )
    mock_readonly_context = Mock(spec=ReadonlyContext)
    expected_headers = {"X-Tenant-ID": "test-tenant"}

    async def header_provider(_context):
      return expected_headers

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        header_provider=header_provider,
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools(readonly_context=mock_readonly_context)

    assert len(tools) == 2
    self.mock_session_manager.create_session.assert_called_once_with(
        headers=expected_headers
    )

  @pytest.mark.asyncio
  async def test_close_success(self):
    """Test successful cleanup."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    await toolset.close()

    self.mock_session_manager.close.assert_called_once()

  @pytest.mark.asyncio
  async def test_close_with_exception(self):
    """Test cleanup when session manager raises exception."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    # Mock close to raise an exception
    self.mock_session_manager.close = AsyncMock(
        side_effect=Exception("Cleanup error")
    )

    # Should not raise exception, should log the warning
    await toolset.close()

  @pytest.mark.asyncio
  async def test_get_tools_with_timeout(self):
    """Test get_tools with timeout."""
    stdio_params = StdioConnectionParams(
        server_params=self.mock_stdio_params, timeout=0.01
    )
    toolset = McpToolset(connection_params=stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    async def long_running_list_tools():
      await asyncio.sleep(0.1)
      return MockListToolsResult([])

    self.mock_session.list_tools = long_running_list_tools

    with pytest.raises(
        ConnectionError, match="Failed to get tools from MCP server."
    ):
      await toolset.get_tools()

  @pytest.mark.asyncio
  async def test_get_tools_retry_decorator(self):
    """Test that get_tools has retry decorator applied."""
    toolset = McpToolset(connection_params=self.mock_stdio_params)

    # Check that the method has the retry decorator
    assert hasattr(toolset.get_tools, "__wrapped__")

  @pytest.mark.asyncio
  async def test_mcp_toolset_with_prefix(self):
    """Test that McpToolset correctly applies the tool_name_prefix."""
    # Mock the connection parameters
    mock_connection_params = MagicMock()
    mock_connection_params.timeout = None

    # Mock the MCPSessionManager and its create_session method
    mock_session_manager = MagicMock()
    mock_session = MagicMock()

    # Mock the list_tools response from the MCP server
    mock_tool1 = MagicMock()
    mock_tool1.name = "tool1"
    mock_tool1.description = "tool 1 desc"
    mock_tool2 = MagicMock()
    mock_tool2.name = "tool2"
    mock_tool2.description = "tool 2 desc"
    list_tools_result = MagicMock()
    list_tools_result.tools = [mock_tool1, mock_tool2]
    mock_session.list_tools = AsyncMock(return_value=list_tools_result)
    mock_session_manager.create_session = AsyncMock(return_value=mock_session)

    # Create an instance of McpToolset with a prefix
    toolset = McpToolset(
        connection_params=mock_connection_params,
        tool_name_prefix="my_prefix",
        use_mcp_resources=True,
    )

    # Replace the internal session manager with our mock
    toolset._mcp_session_manager = mock_session_manager

    # Get the tools from the toolset
    tools = await toolset.get_tools()

    # The get_tools method in McpToolset returns MCPTool objects, which are
    # instances of BaseTool. The prefixing is handled by the BaseToolset,
    # so we need to call get_tools_with_prefix to get the prefixed tools.
    prefixed_tools = await toolset.get_tools_with_prefix()

    # Assert that the tools are prefixed correctly
    assert len(prefixed_tools) == 3
    assert prefixed_tools[0].name == "my_prefix_tool1"
    assert prefixed_tools[1].name == "my_prefix_tool2"
    assert prefixed_tools[2].name == "my_prefix_load_mcp_resource"

    # Assert that the original tools are not modified
    assert tools[0].name == "tool1"
    assert tools[1].name == "tool2"
    assert tools[2].name == "load_mcp_resource"

  def test_init_with_progress_callback(self):
    """Test initialization with progress_callback."""

    async def my_progress_callback(
        progress: float, total: float | None, message: str | None
    ) -> None:
      pass

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        progress_callback=my_progress_callback,
    )

    assert toolset._progress_callback == my_progress_callback

  @pytest.mark.asyncio
  async def test_get_tools_passes_progress_callback_to_mcp_tools(self):
    """Test that get_tools passes progress_callback to created MCPTool instances."""
    progress_updates = []

    async def my_progress_callback(
        progress: float, total: float | None, message: str | None
    ) -> None:
      progress_updates.append((progress, total, message))

    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        progress_callback=my_progress_callback,
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools()

    assert len(tools) == 2
    # Verify each tool has the progress_callback set
    for tool in tools:
      assert tool._progress_callback == my_progress_callback

  def test_init_with_progress_callback_factory(self):
    """Test initialization with a ProgressCallbackFactory."""

    def my_callback_factory(tool_name: str, *, readonly_context=None, **kwargs):
      async def callback(
          progress: float, total: float | None, message: str | None
      ) -> None:
        pass

      return callback

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        progress_callback=my_callback_factory,
    )

    assert toolset._progress_callback == my_callback_factory

  @pytest.mark.asyncio
  async def test_get_tools_passes_factory_to_mcp_tools(self):
    """Test that get_tools passes factory directly to MCPTool instances.

    The factory is resolved at runtime in McpTool._run_async_impl, not at
    tool creation time. This allows the factory to receive ReadonlyContext.
    """

    def my_callback_factory(tool_name: str, *, readonly_context=None, **kwargs):
      async def callback(
          progress: float, total: float | None, message: str | None
      ) -> None:
        pass

      return callback

    mock_tools = [MockMCPTool("tool1"), MockMCPTool("tool2")]
    self.mock_session.list_tools = AsyncMock(
        return_value=MockListToolsResult(mock_tools)
    )

    toolset = McpToolset(
        connection_params=self.mock_stdio_params,
        progress_callback=my_callback_factory,
    )
    toolset._mcp_session_manager = self.mock_session_manager

    tools = await toolset.get_tools()

    assert len(tools) == 2
    # Factory is passed directly to each tool (resolved at runtime)
    for tool in tools:
      assert tool._progress_callback == my_callback_factory

  @pytest.mark.asyncio
  async def test_list_resources(self):
    """Test listing resources."""
    resources = [
        Resource(
            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
        ),
        Resource(
            name="data.json",
            mime_type="application/json",
            uri="file:///data.json",
        ),
    ]
    list_resources_result = ListResourcesResult(resources=resources)
    self.mock_session.list_resources = AsyncMock(
        return_value=list_resources_result
    )

    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    result = await toolset.list_resources()

    assert result == ["file1.txt", "data.json"]
    self.mock_session.list_resources.assert_called_once()

  @pytest.mark.asyncio
  async def test_get_resource_info_success(self):
    """Test getting resource info for an existing resource."""
    resources = [
        Resource(
            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
        ),
        Resource(
            name="data.json",
            mime_type="application/json",
            uri="file:///data.json",
        ),
    ]
    list_resources_result = ListResourcesResult(resources=resources)
    self.mock_session.list_resources = AsyncMock(
        return_value=list_resources_result
    )

    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    result = await toolset.get_resource_info("data.json")

    assert result == {
        "name": "data.json",
        "mime_type": "application/json",
        "uri": "file:///data.json",
    }
    self.mock_session.list_resources.assert_called_once()

  @pytest.mark.asyncio
  async def test_get_resource_info_not_found(self):
    """Test getting resource info for a non-existent resource."""
    resources = [
        Resource(
            name="file1.txt", mime_type="text/plain", uri="file:///file1.txt"
        ),
    ]
    list_resources_result = ListResourcesResult(resources=resources)
    self.mock_session.list_resources = AsyncMock(
        return_value=list_resources_result
    )

    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    with pytest.raises(
        ValueError, match="Resource with name 'other.json' not found."
    ):
      await toolset.get_resource_info("other.json")

  @pytest.mark.parametrize(
      "name,mime_type,content,encoding",
      [
          ("file1.txt", "text/plain", "hello world", None),
          (
              "data.json",
              "application/json",
              '{"key": "value"}',
              None,
          ),
          (
              "file1_b64.txt",
              "text/plain",
              base64.b64encode(b"hello world").decode("ascii"),
              "base64",
          ),
          (
              "data_b64.json",
              "application/json",
              base64.b64encode(b'{"key": "value"}').decode("ascii"),
              "base64",
          ),
          (
              "data.bin",
              "application/octet-stream",
              base64.b64encode(b"\x01\x02\x03").decode("ascii"),
              "base64",
          ),
      ],
  )
  @pytest.mark.asyncio
  async def test_read_resource(self, name, mime_type, content, encoding):
    """Test reading various resource types."""
    uri = f"file:///{name}"
    # Mock list_resources for get_resource_info
    resources = [Resource(name=name, mime_type=mime_type, uri=uri)]
    list_resources_result = ListResourcesResult(resources=resources)
    self.mock_session.list_resources = AsyncMock(
        return_value=list_resources_result
    )

    # Mock read_resource
    if encoding == "base64":
      contents = [
          BlobResourceContents(uri=uri, mimeType=mime_type, blob=content)
      ]
    else:
      contents = [
          TextResourceContents(uri=uri, mimeType=mime_type, text=content)
      ]

    read_resource_result = ReadResourceResult(contents=contents)
    self.mock_session.read_resource = AsyncMock(
        return_value=read_resource_result
    )

    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._mcp_session_manager = self.mock_session_manager

    result = await toolset.read_resource(name)

    assert result == contents
    self.mock_session.list_resources.assert_called_once()
    self.mock_session.read_resource.assert_called_once_with(uri=uri)

  @pytest.mark.asyncio
  async def test_sampling_callback_invoked(self):

    called = {"value": False}

    async def mock_sampling_handler(messages, params=None, context=None):
      called["value"] = True

      assert isinstance(messages, list)
      assert messages[0]["role"] == "user"

      return {
          "model": "test-model",
          "role": "assistant",
          "content": {"type": "text", "text": "sampling response"},
          "stopReason": "endTurn",
      }

    toolset = McpToolset(
        connection_params=StreamableHTTPConnectionParams(
            url="http://localhost:9999",
            timeout=10,
        ),
        sampling_callback=mock_sampling_handler,
    )

    messages = [{"role": "user", "content": {"type": "text", "text": "hello"}}]

    result = await toolset._sampling_callback(messages)

    assert called["value"] is True
    assert result["role"] == "assistant"
    assert result["content"]["text"] == "sampling response"

  @pytest.mark.asyncio
  async def test_get_auth_headers_includes_additional_headers(self):
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.HTTP,
        http=HttpAuth(
            scheme="bearer",
            credentials=HttpCredentials(token="token"),
            additional_headers={"X-API-Key": "secret"},
        ),
    )
    auth_config = AuthConfig(
        auth_scheme=OAuth2(flows={}),
        raw_auth_credential=credential,
    )
    auth_config.exchanged_auth_credential = credential
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    toolset._auth_config = auth_config

    headers = toolset._get_auth_headers()

    assert headers["Authorization"] == "Bearer token"
    assert headers["X-API-Key"] == "secret"

  def test_pickle_mcp_toolset(self):
    toolset = McpToolset(connection_params=self.mock_stdio_params)
    pickled = pickle.dumps(toolset)
    unpickled = pickle.loads(pickled)
    assert unpickled._connection_params == self.mock_stdio_params
    assert unpickled._errlog == sys.stderr