mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Add SSL certificate verification configuration to OpenAPI tools
This change introduces a `verify` parameter to `RestApiTool` and `OpenAPIToolset`. This parameter allows users to configure how SSL certificates are verified when making API calls using the `requests` library. Options include providing a path to a CA bundle, disabling verification, or using a custom `ssl.SSLContext`. New methods `configure_verify` and `configure_verify_all` are added to update this setting after initialization. This is useful for environments with TLS-intercepting proxies. Fixes: https://github.com/google/adk-python/issues/3720 Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 840809727
This commit is contained in:
committed by
Copybara-Service
parent
711df01e73
commit
9d2388a46f
@@ -13,6 +13,7 @@
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
from typing import Dict
|
||||
|
||||
from fastapi.openapi.models import APIKey
|
||||
@@ -137,3 +138,34 @@ def test_openapi_toolset_configure_auth_on_init(openapi_spec: Dict):
|
||||
for tool in toolset._tools:
|
||||
assert tool.auth_scheme == auth_scheme
|
||||
assert tool.auth_credential == auth_credential
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"verify_value", ["/path/to/enterprise-ca-bundle.crt", False]
|
||||
)
|
||||
def test_openapi_toolset_verify_on_init(
|
||||
openapi_spec: Dict[str, Any], verify_value: str | bool
|
||||
):
|
||||
"""Test configuring verify during initialization."""
|
||||
toolset = OpenAPIToolset(
|
||||
spec_dict=openapi_spec,
|
||||
ssl_verify=verify_value,
|
||||
)
|
||||
for tool in toolset._tools:
|
||||
assert tool._ssl_verify == verify_value
|
||||
|
||||
|
||||
def test_openapi_toolset_configure_verify_all(openapi_spec: Dict[str, Any]):
|
||||
"""Test configure_verify_all method."""
|
||||
toolset = OpenAPIToolset(spec_dict=openapi_spec)
|
||||
|
||||
# Initially verify should be None
|
||||
for tool in toolset._tools:
|
||||
assert tool._ssl_verify is None
|
||||
|
||||
# Configure verify for all tools
|
||||
ca_bundle_path = "/path/to/custom-ca.crt"
|
||||
toolset.configure_ssl_verify_all(ca_bundle_path)
|
||||
|
||||
for tool in toolset._tools:
|
||||
assert tool._ssl_verify == ca_bundle_path
|
||||
|
||||
@@ -14,6 +14,8 @@
|
||||
|
||||
|
||||
import json
|
||||
import ssl
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch
|
||||
@@ -48,6 +50,11 @@ class TestRestApiTool:
|
||||
mock_context.request_credential.return_value = {}
|
||||
return mock_context
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ssl_context(self):
|
||||
"""Fixture for a mock ssl.SSLContext."""
|
||||
return mock.create_autospec(ssl.SSLContext)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_operation_parser(self):
|
||||
"""Fixture for a mock OperationParser."""
|
||||
@@ -934,6 +941,101 @@ class TestRestApiTool:
|
||||
assert "param_name" in request_params["params"]
|
||||
assert "empty_param" not in request_params["params"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"verify_input, expected_verify_in_call",
|
||||
[
|
||||
(True, True),
|
||||
(False, False),
|
||||
(
|
||||
"/path/to/enterprise-ca-bundle.crt",
|
||||
"/path/to/enterprise-ca-bundle.crt",
|
||||
),
|
||||
(
|
||||
"USE_SSL_FIXTURE",
|
||||
"USE_SSL_FIXTURE",
|
||||
),
|
||||
(None, None), # None means 'verify' should not be in call_kwargs
|
||||
],
|
||||
)
|
||||
async def test_call_with_verify_options(
|
||||
self,
|
||||
mock_tool_context,
|
||||
sample_endpoint,
|
||||
sample_operation,
|
||||
sample_auth_scheme,
|
||||
sample_auth_credential,
|
||||
mock_ssl_context,
|
||||
verify_input,
|
||||
expected_verify_in_call,
|
||||
):
|
||||
"""Test different values for the 'verify' parameter."""
|
||||
if verify_input == "USE_SSL_FIXTURE":
|
||||
verify_input = mock_ssl_context
|
||||
if expected_verify_in_call == "USE_SSL_FIXTURE":
|
||||
expected_verify_in_call = mock_ssl_context
|
||||
|
||||
mock_response = mock.create_autospec(
|
||||
requests.Response, instance=True, spec_set=True
|
||||
)
|
||||
mock_response.json.return_value = {"result": "success"}
|
||||
|
||||
tool = RestApiTool(
|
||||
name="test_tool",
|
||||
description="Test Tool",
|
||||
endpoint=sample_endpoint,
|
||||
operation=sample_operation,
|
||||
auth_scheme=sample_auth_scheme,
|
||||
auth_credential=sample_auth_credential,
|
||||
ssl_verify=verify_input,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
requests, "request", return_value=mock_response, autospec=True
|
||||
) as mock_request:
|
||||
await tool.call(args={}, tool_context=mock_tool_context)
|
||||
|
||||
assert mock_request.called
|
||||
_, call_kwargs = mock_request.call_args
|
||||
if expected_verify_in_call is None:
|
||||
assert "verify" not in call_kwargs
|
||||
else:
|
||||
assert call_kwargs["verify"] == expected_verify_in_call
|
||||
|
||||
async def test_call_with_configure_verify(
|
||||
self,
|
||||
mock_tool_context,
|
||||
sample_endpoint,
|
||||
sample_operation,
|
||||
sample_auth_scheme,
|
||||
sample_auth_credential,
|
||||
):
|
||||
"""Test that configure_verify updates the verify setting."""
|
||||
mock_response = mock.create_autospec(
|
||||
requests.Response, instance=True, spec_set=True
|
||||
)
|
||||
mock_response.json.return_value = {"result": "success"}
|
||||
|
||||
tool = RestApiTool(
|
||||
name="test_tool",
|
||||
description="Test Tool",
|
||||
endpoint=sample_endpoint,
|
||||
operation=sample_operation,
|
||||
auth_scheme=sample_auth_scheme,
|
||||
auth_credential=sample_auth_credential,
|
||||
)
|
||||
|
||||
ca_bundle_path = "/path/to/custom-ca.crt"
|
||||
tool.configure_ssl_verify(ca_bundle_path)
|
||||
|
||||
with patch.object(
|
||||
requests, "request", return_value=mock_response
|
||||
) as mock_request:
|
||||
await tool.call(args={}, tool_context=mock_tool_context)
|
||||
|
||||
assert mock_request.called
|
||||
call_kwargs = mock_request.call_args[1]
|
||||
assert call_kwargs["verify"] == ca_bundle_path
|
||||
|
||||
|
||||
def test_snake_to_lower_camel():
|
||||
assert snake_to_lower_camel("single") == "single"
|
||||
|
||||
Reference in New Issue
Block a user