Skip to content
31 changes: 31 additions & 0 deletions src/google/adk/tools/mcp_tool/mcp_toolset.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,37 @@ def _get_auth_headers(self) -> Optional[Dict[str, str]]:

return headers

@property
def connection_params(self) -> Union[
StdioServerParameters,
StdioConnectionParams,
SseConnectionParams,
StreamableHTTPConnectionParams,
]:
return self._connection_params

@property
def auth_scheme(self) -> Optional[AuthScheme]:
return self._auth_scheme

@property
def auth_credential(self) -> Optional[AuthCredential]:
return self._auth_credential

@property
def require_confirmation(self) -> Union[bool, Callable[..., bool]]:
return self._require_confirmation

@property
def header_provider(
self,
) -> Optional[Callable[[ReadonlyContext], Dict[str, str]]]:
Comment on lines +251 to +253
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The formatting of this property definition is unconventional. It's more standard to keep self on the same line as the method name. If the line length is an issue, you can break the type hint across multiple lines for better readability.

  def header_provider(self) -> Optional[Callable[[ReadonlyContext], Dict[str, str]]]:

Comment on lines +251 to +253
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The formatting of the header_provider property definition is unconventional. For better readability and consistency with the other properties, the method signature should be on a single line.

  def header_provider(self) -> Optional[Callable[[ReadonlyContext], Dict[str, str]]]:

return self._header_provider

@property
def errlog(self) -> TextIO:
return self._errlog

async def _execute_with_session(
self,
coroutine_func: Callable[[Any], Awaitable[T]],
Expand Down
55 changes: 55 additions & 0 deletions tests/unittests/tools/mcp_tool/test_mcp_toolset.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,61 @@ def test_init_basic(self):
assert toolset._auth_scheme is None
assert toolset._auth_credential is None

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."""
mock_scheme = Mock()
toolset = MCPToolset(
connection_params=self.mock_stdio_params,
auth_scheme=mock_scheme,
)
assert toolset.auth_scheme == mock_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
Comment on lines 85 to 138
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The new tests are a good start, but they don't cover all the newly added properties. The properties require_confirmation and header_provider are not tested. Additionally, auth_scheme is only tested for its default None value, unlike auth_credential which has a separate test for a non-default value. To ensure the new properties are fully validated, please consider adding tests for the missing properties and for the non-default case of auth_scheme.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you also add the additional tests.


def test_init_with_stdio_connection_params(self):
"""Test initialization with StdioConnectionParams."""
stdio_params = StdioConnectionParams(
Expand Down