chore: import upstream snapshot with attribution
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
+705
@@ -0,0 +1,705 @@
|
||||
# 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 json
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
|
||||
import google.auth
|
||||
import pytest
|
||||
import requests
|
||||
from requests import exceptions
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project():
|
||||
return "test-project"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def location():
|
||||
return "us-central1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_name():
|
||||
return "test-connection"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credentials():
|
||||
creds = mock.create_autospec(google.auth.credentials.Credentials)
|
||||
creds.token = "test_token"
|
||||
creds.expired = False
|
||||
return creds
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_request():
|
||||
return mock.create_autospec(google.auth.transport.requests.Request)
|
||||
|
||||
|
||||
class TestConnectionsClient:
|
||||
|
||||
def test_initialization(self, project, location, connection_name):
|
||||
credentials = {"email": "test@example.com"}
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
|
||||
return_value="connectors.googleapis.com",
|
||||
) as mock_get_api_endpoint:
|
||||
client = ConnectionsClient(
|
||||
project, location, connection_name, json.dumps(credentials)
|
||||
)
|
||||
assert client.project == project
|
||||
assert client.location == location
|
||||
assert client.connection == connection_name
|
||||
assert client.connector_url == "https://connectors.googleapis.com"
|
||||
assert client.service_account_json == json.dumps(credentials)
|
||||
assert client.credential_cache is None
|
||||
mock_get_api_endpoint.assert_called_once_with(
|
||||
location,
|
||||
"connectors.googleapis.com",
|
||||
"connectors.mtls.googleapis.com",
|
||||
)
|
||||
|
||||
def test_initialization_mtls_endpoint(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
|
||||
return_value="connectors.mtls.googleapis.com",
|
||||
):
|
||||
client = ConnectionsClient(
|
||||
project, location, connection_name, json.dumps(credentials)
|
||||
)
|
||||
assert client.connector_url == "https://connectors.mtls.googleapis.com"
|
||||
|
||||
def test_execute_api_call_success(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_response.json.return_value = {"data": "test"}
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
client, "_get_access_token", return_value=mock_credentials.token
|
||||
),
|
||||
mock.patch("requests.get", return_value=mock_response),
|
||||
):
|
||||
response = client._execute_api_call("https://test.url")
|
||||
assert response.json() == {"data": "test"}
|
||||
requests.get.assert_called_once_with(
|
||||
"https://test.url",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {mock_credentials.token}",
|
||||
},
|
||||
)
|
||||
|
||||
def test_execute_api_call_credential_error(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
with mock.patch.object(
|
||||
client,
|
||||
"_get_access_token",
|
||||
side_effect=google.auth.exceptions.DefaultCredentialsError("Test"),
|
||||
):
|
||||
with pytest.raises(PermissionError, match="Credentials error: Test"):
|
||||
client._execute_api_call("https://test.url")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code, response_text",
|
||||
[(404, "Not Found"), (400, "Bad Request")],
|
||||
)
|
||||
def test_execute_api_call_request_error_not_found_or_bad_request(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
connection_name,
|
||||
mock_credentials,
|
||||
status_code,
|
||||
response_text,
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.raise_for_status.side_effect = exceptions.HTTPError(
|
||||
f"HTTP error {status_code}: {response_text}"
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
client, "_get_access_token", return_value=mock_credentials.token
|
||||
),
|
||||
mock.patch("requests.get", return_value=mock_response),
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="Invalid request. Please check the provided"
|
||||
):
|
||||
client._execute_api_call("https://test.url")
|
||||
|
||||
def test_execute_api_call_other_request_error(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.raise_for_status.side_effect = exceptions.HTTPError(
|
||||
"Internal Server Error"
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
client, "_get_access_token", return_value=mock_credentials.token
|
||||
),
|
||||
mock.patch("requests.get", return_value=mock_response),
|
||||
):
|
||||
with pytest.raises(ValueError, match="Request error: "):
|
||||
client._execute_api_call("https://test.url")
|
||||
|
||||
def test_execute_api_call_unexpected_error(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
with (
|
||||
mock.patch.object(
|
||||
client, "_get_access_token", return_value=mock_credentials.token
|
||||
),
|
||||
mock.patch(
|
||||
"requests.get", side_effect=Exception("Something went wrong")
|
||||
),
|
||||
):
|
||||
with pytest.raises(
|
||||
Exception, match="An unexpected error occurred: Something went wrong"
|
||||
):
|
||||
client._execute_api_call("https://test.url")
|
||||
|
||||
def test_get_connection_details_success_with_host(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"name": "test-connection",
|
||||
"serviceDirectory": "test_service",
|
||||
"host": "test.host",
|
||||
"tlsServiceDirectory": "tls_test_service",
|
||||
"authOverrideEnabled": True,
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", return_value=mock_response
|
||||
):
|
||||
details = client.get_connection_details()
|
||||
assert details == {
|
||||
"name": "test-connection",
|
||||
"serviceName": "tls_test_service",
|
||||
"host": "test.host",
|
||||
"authOverrideEnabled": True,
|
||||
}
|
||||
|
||||
def test_get_connection_details_success_without_host(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"name": "test-connection",
|
||||
"serviceDirectory": "test_service",
|
||||
"authOverrideEnabled": False,
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", return_value=mock_response
|
||||
):
|
||||
details = client.get_connection_details()
|
||||
assert details == {
|
||||
"name": "test-connection",
|
||||
"serviceName": "test_service",
|
||||
"host": "",
|
||||
"authOverrideEnabled": False,
|
||||
}
|
||||
|
||||
def test_get_connection_details_without_name(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"serviceDirectory": "test_service",
|
||||
"authOverrideEnabled": False,
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", return_value=mock_response
|
||||
):
|
||||
details = client.get_connection_details()
|
||||
assert details == {
|
||||
"name": "",
|
||||
"serviceName": "test_service",
|
||||
"host": "",
|
||||
"authOverrideEnabled": False,
|
||||
}
|
||||
|
||||
def test_get_connection_details_error(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", side_effect=ValueError("Request error")
|
||||
):
|
||||
with pytest.raises(ValueError, match="Request error"):
|
||||
client.get_connection_details()
|
||||
|
||||
def test_get_entity_schema_and_operations_success(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
|
||||
return_value="connectors.googleapis.com",
|
||||
):
|
||||
client = ConnectionsClient(
|
||||
project, location, connection_name, credentials
|
||||
)
|
||||
mock_execute_response_initial = mock.MagicMock()
|
||||
mock_execute_response_initial.status_code = 200
|
||||
mock_execute_response_initial.json.return_value = {
|
||||
"name": "operations/test_op"
|
||||
}
|
||||
|
||||
mock_execute_response_poll_done = mock.MagicMock()
|
||||
mock_execute_response_poll_done.status_code = 200
|
||||
mock_execute_response_poll_done.json.return_value = {
|
||||
"done": True,
|
||||
"response": {
|
||||
"jsonSchema": {"type": "object"},
|
||||
"operations": ["LIST", "GET"],
|
||||
},
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
client,
|
||||
"_execute_api_call",
|
||||
side_effect=[
|
||||
mock_execute_response_initial,
|
||||
mock_execute_response_poll_done,
|
||||
],
|
||||
):
|
||||
schema, operations = client.get_entity_schema_and_operations("entity1")
|
||||
assert schema == {"type": "object"}
|
||||
assert operations == ["LIST", "GET"]
|
||||
assert (
|
||||
mock.call(
|
||||
f"https://connectors.googleapis.com/v1/projects/{project}/locations/{location}/connections/{connection_name}/connectionSchemaMetadata:getEntityType?entityId=entity1"
|
||||
)
|
||||
in client._execute_api_call.mock_calls
|
||||
)
|
||||
assert (
|
||||
mock.call(f"https://connectors.googleapis.com/v1/operations/test_op")
|
||||
in client._execute_api_call.mock_calls
|
||||
)
|
||||
|
||||
def test_get_entity_schema_and_operations_no_operation_id(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_execute_response = mock.MagicMock()
|
||||
mock_execute_response.status_code = 200
|
||||
mock_execute_response.json.return_value = {}
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", return_value=mock_execute_response
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Failed to get entity schema and operations for entity: entity1"
|
||||
),
|
||||
):
|
||||
client.get_entity_schema_and_operations("entity1")
|
||||
|
||||
def test_get_entity_schema_and_operations_execute_api_call_error(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", side_effect=ValueError("Request error")
|
||||
):
|
||||
with pytest.raises(ValueError, match="Request error"):
|
||||
client.get_entity_schema_and_operations("entity1")
|
||||
|
||||
def test_get_action_schema_success(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
|
||||
return_value="connectors.googleapis.com",
|
||||
):
|
||||
client = ConnectionsClient(
|
||||
project, location, connection_name, credentials
|
||||
)
|
||||
mock_execute_response_initial = mock.MagicMock()
|
||||
mock_execute_response_initial.status_code = 200
|
||||
mock_execute_response_initial.json.return_value = {
|
||||
"name": "operations/test_op"
|
||||
}
|
||||
|
||||
mock_execute_response_poll_done = mock.MagicMock()
|
||||
mock_execute_response_poll_done.status_code = 200
|
||||
mock_execute_response_poll_done.json.return_value = {
|
||||
"done": True,
|
||||
"response": {
|
||||
"inputJsonSchema": {
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string"}},
|
||||
},
|
||||
"outputJsonSchema": {
|
||||
"type": "object",
|
||||
"properties": {"output": {"type": "string"}},
|
||||
},
|
||||
"description": "Test Action Description",
|
||||
"displayName": "TestAction",
|
||||
},
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
client,
|
||||
"_execute_api_call",
|
||||
side_effect=[
|
||||
mock_execute_response_initial,
|
||||
mock_execute_response_poll_done,
|
||||
],
|
||||
):
|
||||
schema = client.get_action_schema("action1")
|
||||
assert schema == {
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string"}},
|
||||
},
|
||||
"outputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"output": {"type": "string"}},
|
||||
},
|
||||
"description": "Test Action Description",
|
||||
"displayName": "TestAction",
|
||||
}
|
||||
assert (
|
||||
mock.call(
|
||||
f"https://connectors.googleapis.com/v1/projects/{project}/locations/{location}/connections/{connection_name}/connectionSchemaMetadata:getAction?actionId=action1"
|
||||
)
|
||||
in client._execute_api_call.mock_calls
|
||||
)
|
||||
assert (
|
||||
mock.call(f"https://connectors.googleapis.com/v1/operations/test_op")
|
||||
in client._execute_api_call.mock_calls
|
||||
)
|
||||
|
||||
def test_get_action_schema_no_operation_id(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
mock_execute_response = mock.MagicMock()
|
||||
mock_execute_response.status_code = 200
|
||||
mock_execute_response.json.return_value = {}
|
||||
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", return_value=mock_execute_response
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="Failed to get action schema for action: action1"
|
||||
):
|
||||
client.get_action_schema("action1")
|
||||
|
||||
def test_get_action_schema_execute_api_call_error(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
with mock.patch.object(
|
||||
client, "_execute_api_call", side_effect=ValueError("Request error")
|
||||
):
|
||||
with pytest.raises(ValueError, match="Request error"):
|
||||
client.get_action_schema("action1")
|
||||
|
||||
def test_get_connector_base_spec(self):
|
||||
spec = ConnectionsClient.get_connector_base_spec()
|
||||
assert "openapi" in spec
|
||||
assert spec["info"]["title"] == "ExecuteConnection"
|
||||
assert "components" in spec
|
||||
assert "schemas" in spec["components"]
|
||||
assert "operation" in spec["components"]["schemas"]
|
||||
|
||||
def test_get_action_operation(self):
|
||||
operation = ConnectionsClient.get_action_operation(
|
||||
"TestAction", "EXECUTE_ACTION", "TestActionDisplayName", "test_tool"
|
||||
)
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "TestActionDisplayName"
|
||||
assert "operationId" in operation["post"]
|
||||
assert operation["post"]["operationId"] == "test_tool_TestActionDisplayName"
|
||||
|
||||
def test_list_operation(self):
|
||||
operation = ConnectionsClient.list_operation(
|
||||
"Entity1", '{"type": "object"}', "test_tool"
|
||||
)
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "List Entity1"
|
||||
assert "operationId" in operation["post"]
|
||||
assert operation["post"]["operationId"] == "test_tool_list_Entity1"
|
||||
|
||||
def test_get_operation_static(self):
|
||||
operation = ConnectionsClient.get_operation(
|
||||
"Entity1", '{"type": "object"}', "test_tool"
|
||||
)
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "Get Entity1"
|
||||
assert "operationId" in operation["post"]
|
||||
assert operation["post"]["operationId"] == "test_tool_get_Entity1"
|
||||
|
||||
def test_create_operation(self):
|
||||
operation = ConnectionsClient.create_operation("Entity1", "test_tool")
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "Creates a new Entity1"
|
||||
assert "operationId" in operation["post"]
|
||||
assert operation["post"]["operationId"] == "test_tool_create_Entity1"
|
||||
|
||||
def test_update_operation(self):
|
||||
operation = ConnectionsClient.update_operation("Entity1", "test_tool")
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "Updates the Entity1"
|
||||
assert "operationId" in operation["post"]
|
||||
assert operation["post"]["operationId"] == "test_tool_update_Entity1"
|
||||
|
||||
def test_delete_operation(self):
|
||||
operation = ConnectionsClient.delete_operation("Entity1", "test_tool")
|
||||
assert "post" in operation
|
||||
assert operation["post"]["summary"] == "Delete the Entity1"
|
||||
assert operation["post"]["operationId"] == "test_tool_delete_Entity1"
|
||||
|
||||
def test_create_operation_request(self):
|
||||
schema = ConnectionsClient.create_operation_request("Entity1")
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "connectorInputPayload" in schema["properties"]
|
||||
|
||||
def test_update_operation_request(self):
|
||||
schema = ConnectionsClient.update_operation_request("Entity1")
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "entityId" in schema["properties"]
|
||||
assert "filterClause" in schema["properties"]
|
||||
|
||||
def test_get_operation_request_static(self):
|
||||
schema = ConnectionsClient.get_operation_request()
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "entityId" in schema["properties"]
|
||||
|
||||
def test_delete_operation_request(self):
|
||||
schema = ConnectionsClient.delete_operation_request()
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "entityId" in schema["properties"]
|
||||
assert "filterClause" in schema["properties"]
|
||||
|
||||
def test_list_operation_request(self):
|
||||
schema = ConnectionsClient.list_operation_request()
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "filterClause" in schema["properties"]
|
||||
assert "sortByColumns" in schema["properties"]
|
||||
|
||||
def test_action_request(self):
|
||||
schema = ConnectionsClient.action_request("TestAction")
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "connectorInputPayload" in schema["properties"]
|
||||
|
||||
def test_action_response(self):
|
||||
schema = ConnectionsClient.action_response("TestAction")
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "connectorOutputPayload" in schema["properties"]
|
||||
|
||||
def test_execute_custom_query_request(self):
|
||||
schema = ConnectionsClient.execute_custom_query_request()
|
||||
assert "type" in schema
|
||||
assert schema["type"] == "object"
|
||||
assert "properties" in schema
|
||||
assert "query" in schema["properties"]
|
||||
|
||||
def test_connector_payload(self):
|
||||
client = ConnectionsClient("test-project", "us-central1", "test-connection")
|
||||
schema = client.connector_payload(
|
||||
json_schema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"input": {
|
||||
"type": ["null", "string"],
|
||||
"description": "description",
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
assert schema == {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"input": {
|
||||
"type": "string",
|
||||
"nullable": True,
|
||||
"description": "description",
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
def test_get_access_token_uses_cached_token(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
credentials = {"email": "test@example.com"}
|
||||
client = ConnectionsClient(project, location, connection_name, credentials)
|
||||
client.credential_cache = mock_credentials
|
||||
token = client._get_access_token()
|
||||
assert token == "test_token"
|
||||
|
||||
def test_get_access_token_with_service_account_credentials(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
service_account_json = json.dumps({
|
||||
"client_email": "test@example.com",
|
||||
"private_key": "test_key",
|
||||
})
|
||||
client = ConnectionsClient(
|
||||
project, location, connection_name, service_account_json
|
||||
)
|
||||
mock_creds = mock.create_autospec(google.oauth2.service_account.Credentials)
|
||||
mock_creds.token = "sa_token"
|
||||
mock_creds.expired = False
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_info",
|
||||
return_value=mock_creds,
|
||||
),
|
||||
mock.patch.object(mock_creds, "refresh", return_value=None),
|
||||
):
|
||||
token = client._get_access_token()
|
||||
assert token == "sa_token"
|
||||
google.oauth2.service_account.Credentials.from_service_account_info.assert_called_once_with(
|
||||
json.loads(service_account_json),
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
mock_creds.refresh.assert_called_once()
|
||||
|
||||
def test_get_access_token_with_default_credentials(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
client = ConnectionsClient(project, location, connection_name, None)
|
||||
with (
|
||||
mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client.default_service_credential",
|
||||
return_value=(mock_credentials, "test_project_id"),
|
||||
) as mock_default_service_credential,
|
||||
mock.patch.object(mock_credentials, "refresh", return_value=None),
|
||||
):
|
||||
token = client._get_access_token()
|
||||
assert token == "test_token"
|
||||
# Verify default_service_credential is called with the correct scopes parameter
|
||||
mock_default_service_credential.assert_called_once_with(
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
||||
)
|
||||
|
||||
def test_get_access_token_no_valid_credentials(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
client = ConnectionsClient(project, location, connection_name, None)
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client.default_service_credential",
|
||||
return_value=(None, None),
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Please provide a service account that has the required"
|
||||
" permissions"
|
||||
),
|
||||
):
|
||||
client._get_access_token()
|
||||
|
||||
def test_get_access_token_default_credentials_error(
|
||||
self, project, location, connection_name
|
||||
):
|
||||
client = ConnectionsClient(project, location, connection_name, None)
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client.default_service_credential",
|
||||
side_effect=google.auth.exceptions.DefaultCredentialsError(
|
||||
"ADC not found"
|
||||
),
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Please provide a service account that has the required"
|
||||
" permissions"
|
||||
),
|
||||
):
|
||||
client._get_access_token()
|
||||
|
||||
def test_get_access_token_refreshes_expired_token(
|
||||
self, project, location, connection_name, mock_credentials
|
||||
):
|
||||
client = ConnectionsClient(project, location, connection_name, None)
|
||||
mock_credentials.expired = True
|
||||
mock_credentials.token = "old_token"
|
||||
mock_credentials.refresh.return_value = None
|
||||
|
||||
client.credential_cache = mock_credentials
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.connections_client.default_service_credential",
|
||||
return_value=(mock_credentials, "test_project_id"),
|
||||
):
|
||||
# Mock the refresh method directly on the instance within the context
|
||||
with mock.patch.object(mock_credentials, "refresh") as mock_refresh:
|
||||
mock_credentials.token = "new_token" # Set the expected new token
|
||||
token = client._get_access_token()
|
||||
assert token == "new_token"
|
||||
mock_refresh.assert_called_once()
|
||||
+756
@@ -0,0 +1,756 @@
|
||||
# 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 json
|
||||
import re
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.tools.application_integration_tool.clients import integration_client
|
||||
from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
|
||||
from google.adk.tools.application_integration_tool.clients.integration_client import IntegrationClient
|
||||
import google.auth
|
||||
import google.auth.transport.requests
|
||||
from google.auth.transport.requests import Request
|
||||
from google.oauth2 import service_account
|
||||
import pytest
|
||||
import requests
|
||||
from requests import exceptions
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project():
|
||||
return "test-project"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def location():
|
||||
return "us-central1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_name():
|
||||
return "test-integration"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def triggers():
|
||||
return ["test-trigger", "test-trigger2"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_name():
|
||||
return "test-connection"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credentials():
|
||||
creds = mock.create_autospec(google.auth.credentials.Credentials)
|
||||
creds.token = "test_token"
|
||||
return creds
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_auth_request():
|
||||
return mock.create_autospec(Request)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_connections_client():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client.ConnectionsClient"
|
||||
) as mock_client:
|
||||
mock_instance = mock.create_autospec(ConnectionsClient)
|
||||
mock_client.return_value = mock_instance
|
||||
yield mock_client
|
||||
|
||||
|
||||
class TestIntegrationClient:
|
||||
|
||||
def test_initialization(
|
||||
self, project, location, integration_name, triggers, connection_name
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations={"entity": ["LIST"]},
|
||||
actions=["action1"],
|
||||
service_account_json=json.dumps({"email": "test@example.com"}),
|
||||
)
|
||||
assert client.project == project
|
||||
assert client.location == location
|
||||
assert client.integration == integration_name
|
||||
assert client.triggers == triggers
|
||||
assert client.connection == connection_name
|
||||
assert client.entity_operations == {"entity": ["LIST"]}
|
||||
assert client.actions == ["action1"]
|
||||
assert client.service_account_json == json.dumps(
|
||||
{"email": "test@example.com"}
|
||||
)
|
||||
assert client.credential_cache is None
|
||||
|
||||
def test_get_openapi_spec_for_integration_success(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_credentials,
|
||||
mock_connections_client,
|
||||
):
|
||||
mock_credentials.quota_project_id = "quota-project"
|
||||
mock_credentials.expired = False
|
||||
expected_spec = {"openapi": "3.0.0", "info": {"title": "Test Integration"}}
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"openApiSpec": json.dumps(expected_spec)}
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
integration_client,
|
||||
"default_service_credential",
|
||||
return_value=(mock_credentials, project),
|
||||
),
|
||||
mock.patch.object(mock_credentials, "refresh", return_value=None),
|
||||
mock.patch.object(requests, "post", return_value=mock_response),
|
||||
mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client._mtls_utils.get_api_endpoint",
|
||||
return_value=f"{location}-integrations.googleapis.com",
|
||||
) as mock_get_api_endpoint,
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
spec = client.get_openapi_spec_for_integration()
|
||||
assert spec == expected_spec
|
||||
mock_get_api_endpoint.assert_called_once_with(
|
||||
location,
|
||||
"{location}-integrations.googleapis.com",
|
||||
"{location}-integrations.mtls.googleapis.com",
|
||||
)
|
||||
requests.post.assert_called_once_with(
|
||||
f"https://{location}-integrations.googleapis.com/v1/projects/{project}/locations/{location}:generateOpenApiSpec",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {mock_credentials.token}",
|
||||
"x-goog-user-project": "quota-project",
|
||||
},
|
||||
json={
|
||||
"apiTriggerResources": [{
|
||||
"integrationResource": integration_name,
|
||||
"triggerId": triggers,
|
||||
}],
|
||||
"fileFormat": "JSON",
|
||||
},
|
||||
)
|
||||
|
||||
def test_get_openapi_spec_for_integration_success_mtls(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_credentials,
|
||||
mock_connections_client,
|
||||
):
|
||||
mock_credentials.quota_project_id = "quota-project"
|
||||
mock_credentials.expired = False
|
||||
expected_spec = {"openapi": "3.0.0", "info": {"title": "Test Integration"}}
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"openApiSpec": json.dumps(expected_spec)}
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
integration_client,
|
||||
"default_service_credential",
|
||||
return_value=(mock_credentials, project),
|
||||
),
|
||||
mock.patch.object(mock_credentials, "refresh", return_value=None),
|
||||
mock.patch.object(requests, "post", return_value=mock_response),
|
||||
mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client._mtls_utils.get_api_endpoint",
|
||||
return_value=f"{location}-integrations.mtls.googleapis.com",
|
||||
),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
client.get_openapi_spec_for_integration()
|
||||
requests.post.assert_called_once_with(
|
||||
f"https://{location}-integrations.mtls.googleapis.com/v1/projects/{project}/locations/{location}:generateOpenApiSpec",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {mock_credentials.token}",
|
||||
"x-goog-user-project": "quota-project",
|
||||
},
|
||||
json={
|
||||
"apiTriggerResources": [{
|
||||
"integrationResource": integration_name,
|
||||
"triggerId": triggers,
|
||||
}],
|
||||
"fileFormat": "JSON",
|
||||
},
|
||||
)
|
||||
|
||||
def test_get_openapi_spec_for_integration_credential_error(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_connections_client,
|
||||
):
|
||||
with mock.patch.object(
|
||||
IntegrationClient,
|
||||
"_get_access_token",
|
||||
side_effect=ValueError(
|
||||
"Please provide a service account that has the required permissions"
|
||||
" to access the connection."
|
||||
),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
Exception,
|
||||
match=(
|
||||
"An unexpected error occurred: Please provide a service account"
|
||||
" that has the required permissions to access the connection."
|
||||
),
|
||||
):
|
||||
client.get_openapi_spec_for_integration()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code, response_text",
|
||||
[(404, "Not Found"), (400, "Bad Request"), (404, ""), (400, "")],
|
||||
)
|
||||
def test_get_openapi_spec_for_integration_request_error_not_found_or_bad_request(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_credentials,
|
||||
status_code,
|
||||
response_text,
|
||||
mock_connections_client,
|
||||
):
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = status_code
|
||||
mock_response.raise_for_status.side_effect = exceptions.HTTPError(
|
||||
f"HTTP error {status_code}: {response_text}"
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
IntegrationClient,
|
||||
"_get_access_token",
|
||||
return_value=mock_credentials.token,
|
||||
),
|
||||
mock.patch("requests.post", return_value=mock_response),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"Invalid request\. Please check the provided values of"
|
||||
rf" project\({project}\), location\({location}\),"
|
||||
rf" integration\({integration_name}\)."
|
||||
),
|
||||
):
|
||||
client.get_openapi_spec_for_integration()
|
||||
|
||||
def test_get_openapi_spec_for_integration_other_request_error(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_credentials,
|
||||
mock_connections_client,
|
||||
):
|
||||
mock_response = mock.MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.raise_for_status.side_effect = exceptions.HTTPError(
|
||||
"Internal Server Error"
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(
|
||||
IntegrationClient,
|
||||
"_get_access_token",
|
||||
return_value=mock_credentials.token,
|
||||
),
|
||||
mock.patch("requests.post", return_value=mock_response),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match="Request error: "):
|
||||
client.get_openapi_spec_for_integration()
|
||||
|
||||
def test_get_openapi_spec_for_integration_unexpected_error(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
mock_credentials,
|
||||
mock_connections_client,
|
||||
):
|
||||
with (
|
||||
mock.patch.object(
|
||||
IntegrationClient,
|
||||
"_get_access_token",
|
||||
return_value=mock_credentials.token,
|
||||
),
|
||||
mock.patch(
|
||||
"requests.post", side_effect=Exception("Something went wrong")
|
||||
),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=None,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
Exception, match="An unexpected error occurred: Something went wrong"
|
||||
):
|
||||
client.get_openapi_spec_for_integration()
|
||||
|
||||
def test_get_openapi_spec_for_connection_no_entity_operations_or_actions(
|
||||
self, project, location, connection_name, mock_connections_client
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=None,
|
||||
triggers=None,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"No entity operations or actions provided. Please provide at least"
|
||||
" one of them."
|
||||
),
|
||||
):
|
||||
client.get_openapi_spec_for_connection()
|
||||
|
||||
def test_get_openapi_spec_for_connection_with_entity_operations(
|
||||
self, project, location, connection_name, mock_connections_client
|
||||
):
|
||||
entity_operations = {"entity1": ["LIST", "GET"]}
|
||||
|
||||
mock_connections_client_instance = mock_connections_client.return_value
|
||||
mock_connections_client_instance.get_connector_base_spec.return_value = {
|
||||
"components": {"schemas": {}},
|
||||
"paths": {},
|
||||
}
|
||||
mock_connections_client_instance.get_entity_schema_and_operations.return_value = (
|
||||
{"type": "object", "properties": {"id": {"type": "string"}}},
|
||||
["LIST", "GET"],
|
||||
)
|
||||
mock_connections_client_instance.connector_payload.return_value = {
|
||||
"type": "object"
|
||||
}
|
||||
mock_connections_client_instance.list_operation.return_value = {"get": {}}
|
||||
mock_connections_client_instance.list_operation_request.return_value = {
|
||||
"type": "object"
|
||||
}
|
||||
mock_connections_client_instance.get_operation.return_value = {"get": {}}
|
||||
mock_connections_client_instance.get_operation_request.return_value = {
|
||||
"type": "object"
|
||||
}
|
||||
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=None,
|
||||
triggers=None,
|
||||
connection=connection_name,
|
||||
entity_operations=entity_operations,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
spec = client.get_openapi_spec_for_connection()
|
||||
assert "paths" in spec
|
||||
assert (
|
||||
f"/v2/projects/{project}/locations/{location}/integrations/ExecuteConnection:execute?triggerId=api_trigger/ExecuteConnection#list_entity1"
|
||||
in spec["paths"]
|
||||
)
|
||||
assert (
|
||||
f"/v2/projects/{project}/locations/{location}/integrations/ExecuteConnection:execute?triggerId=api_trigger/ExecuteConnection#get_entity1"
|
||||
in spec["paths"]
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_connections_client_instance.get_connector_base_spec.assert_called_once()
|
||||
mock_connections_client_instance.get_entity_schema_and_operations.assert_any_call(
|
||||
"entity1"
|
||||
)
|
||||
mock_connections_client_instance.connector_payload.assert_any_call(
|
||||
{"type": "object", "properties": {"id": {"type": "string"}}}
|
||||
)
|
||||
mock_connections_client_instance.list_operation.assert_called_once()
|
||||
mock_connections_client_instance.get_operation.assert_called_once()
|
||||
|
||||
def test_get_openapi_spec_for_connection_with_actions(
|
||||
self, project, location, connection_name, mock_connections_client
|
||||
):
|
||||
actions = ["TestAction"]
|
||||
mock_connections_client_instance = (
|
||||
mock_connections_client.return_value
|
||||
) # Corrected line
|
||||
mock_connections_client_instance.get_connector_base_spec.return_value = {
|
||||
"components": {"schemas": {}},
|
||||
"paths": {},
|
||||
}
|
||||
mock_connections_client_instance.get_action_schema.return_value = {
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string"}},
|
||||
},
|
||||
"outputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"output": {"type": "string"}},
|
||||
},
|
||||
"displayName": "TestAction",
|
||||
}
|
||||
mock_connections_client_instance.connector_payload.side_effect = [
|
||||
{"type": "object"},
|
||||
{"type": "object"},
|
||||
]
|
||||
mock_connections_client_instance.action_request.return_value = {
|
||||
"type": "object"
|
||||
}
|
||||
mock_connections_client_instance.action_response.return_value = {
|
||||
"type": "object"
|
||||
}
|
||||
mock_connections_client_instance.get_action_operation.return_value = {
|
||||
"post": {}
|
||||
}
|
||||
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=None,
|
||||
triggers=None,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=actions,
|
||||
service_account_json=None,
|
||||
)
|
||||
spec = client.get_openapi_spec_for_connection()
|
||||
assert "paths" in spec
|
||||
assert (
|
||||
f"/v2/projects/{project}/locations/{location}/integrations/ExecuteConnection:execute?triggerId=api_trigger/ExecuteConnection#TestAction"
|
||||
in spec["paths"]
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_connections_client_instance.get_connector_base_spec.assert_called_once()
|
||||
mock_connections_client_instance.get_action_schema.assert_called_once_with(
|
||||
"TestAction"
|
||||
)
|
||||
mock_connections_client_instance.connector_payload.assert_any_call(
|
||||
{"type": "object", "properties": {"input": {"type": "string"}}}
|
||||
)
|
||||
mock_connections_client_instance.connector_payload.assert_any_call(
|
||||
{"type": "object", "properties": {"output": {"type": "string"}}}
|
||||
)
|
||||
mock_connections_client_instance.action_request.assert_called_once_with(
|
||||
"TestAction"
|
||||
)
|
||||
mock_connections_client_instance.action_response.assert_called_once_with(
|
||||
"TestAction"
|
||||
)
|
||||
mock_connections_client_instance.get_action_operation.assert_called_once()
|
||||
|
||||
def test_get_openapi_spec_for_connection_invalid_operation(
|
||||
self, project, location, connection_name, mock_connections_client
|
||||
):
|
||||
entity_operations = {"entity1": ["INVALID"]}
|
||||
mock_connections_client_instance = mock_connections_client.return_value
|
||||
mock_connections_client_instance.get_connector_base_spec.return_value = {
|
||||
"components": {"schemas": {}},
|
||||
"paths": {},
|
||||
}
|
||||
mock_connections_client_instance.get_entity_schema_and_operations.return_value = (
|
||||
{"type": "object", "properties": {"id": {"type": "string"}}},
|
||||
["LIST", "GET"],
|
||||
)
|
||||
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=None,
|
||||
triggers=None,
|
||||
connection=connection_name,
|
||||
entity_operations=entity_operations,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError, match="Invalid operation: INVALID for entity: entity1"
|
||||
):
|
||||
client.get_openapi_spec_for_connection()
|
||||
|
||||
def test_get_access_token_with_service_account_json(
|
||||
self, project, location, integration_name, triggers, connection_name
|
||||
):
|
||||
service_account_json = json.dumps({
|
||||
"client_email": "test@example.com",
|
||||
"private_key": "test_key",
|
||||
})
|
||||
mock_creds = mock.create_autospec(service_account.Credentials)
|
||||
mock_creds.token = "sa_token"
|
||||
mock_creds.expired = False
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_info",
|
||||
return_value=mock_creds,
|
||||
),
|
||||
mock.patch.object(mock_creds, "refresh", return_value=None),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=service_account_json,
|
||||
)
|
||||
token = client._get_access_token()
|
||||
assert token == "sa_token"
|
||||
service_account.Credentials.from_service_account_info.assert_called_once_with(
|
||||
json.loads(service_account_json),
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
mock_creds.refresh.assert_called_once()
|
||||
|
||||
def test_get_access_token_with_default_credentials(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
connection_name,
|
||||
mock_credentials,
|
||||
):
|
||||
mock_credentials.expired = False
|
||||
with (
|
||||
mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client.default_service_credential",
|
||||
return_value=(mock_credentials, "test_project_id"),
|
||||
) as mock_default_service_credential,
|
||||
mock.patch.object(mock_credentials, "refresh", return_value=None),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
token = client._get_access_token()
|
||||
assert token == "test_token"
|
||||
# Verify default_service_credential is called with the correct scopes parameter
|
||||
mock_default_service_credential.assert_called_once_with(
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"]
|
||||
)
|
||||
|
||||
def test_get_access_token_no_valid_credentials(
|
||||
self, project, location, integration_name, triggers, connection_name
|
||||
):
|
||||
with (
|
||||
mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client.default_service_credential",
|
||||
return_value=(None, None),
|
||||
),
|
||||
mock.patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_info",
|
||||
return_value=None,
|
||||
),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
try:
|
||||
client._get_access_token()
|
||||
assert False, "ValueError was not raised" # Explicitly fail if no error
|
||||
except ValueError as e:
|
||||
assert (
|
||||
"Please provide a service account that has the required permissions"
|
||||
" to access the connection."
|
||||
in str(e)
|
||||
)
|
||||
|
||||
def test_get_access_token_default_credentials_error(
|
||||
self, project, location, integration_name, triggers, connection_name
|
||||
):
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client.default_service_credential",
|
||||
side_effect=google.auth.exceptions.DefaultCredentialsError(
|
||||
"ADC not found"
|
||||
),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Please provide a service account that has the required"
|
||||
" permissions to access the connection."
|
||||
),
|
||||
):
|
||||
client._get_access_token()
|
||||
|
||||
def test_get_access_token_uses_cached_token(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
connection_name,
|
||||
mock_credentials,
|
||||
):
|
||||
mock_credentials.token = "cached_token"
|
||||
mock_credentials.expired = False
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
client.credential_cache = mock_credentials # Simulate a cached credential
|
||||
with (
|
||||
mock.patch("google.auth.default") as mock_default,
|
||||
mock.patch(
|
||||
"google.oauth2.service_account.Credentials.from_service_account_info"
|
||||
) as mock_sa,
|
||||
):
|
||||
token = client._get_access_token()
|
||||
assert token == "cached_token"
|
||||
mock_default.assert_not_called()
|
||||
mock_sa.assert_not_called()
|
||||
|
||||
def test_get_access_token_refreshes_expired_token(
|
||||
self,
|
||||
project,
|
||||
location,
|
||||
integration_name,
|
||||
triggers,
|
||||
connection_name,
|
||||
mock_credentials,
|
||||
):
|
||||
mock_credentials = mock.create_autospec(google.auth.credentials.Credentials)
|
||||
mock_credentials.token = "old_token"
|
||||
mock_credentials.expired = True
|
||||
mock_credentials.refresh.return_value = None
|
||||
mock_credentials.token = "new_token" # Simulate token refresh
|
||||
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.clients.integration_client.default_service_credential",
|
||||
return_value=(mock_credentials, "test_project_id"),
|
||||
):
|
||||
client = IntegrationClient(
|
||||
project=project,
|
||||
location=location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
connection=connection_name,
|
||||
entity_operations=None,
|
||||
actions=None,
|
||||
service_account_json=None,
|
||||
)
|
||||
client.credential_cache = mock_credentials
|
||||
token = client._get_access_token()
|
||||
assert token == "new_token"
|
||||
mock_credentials.refresh.assert_called_once()
|
||||
+740
@@ -0,0 +1,740 @@
|
||||
# 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 json
|
||||
from unittest import mock
|
||||
|
||||
from fastapi.openapi.models import Operation
|
||||
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 OAuth2Auth
|
||||
from google.adk.tools.application_integration_tool.application_integration_toolset import ApplicationIntegrationToolset
|
||||
from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool
|
||||
from google.adk.tools.openapi_tool.auth.auth_helpers import dict_to_auth_scheme
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser import rest_api_tool
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_spec_parser import OperationEndpoint
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_spec_parser import ParsedOperation
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_integration_client():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.IntegrationClient"
|
||||
) as mock_client:
|
||||
yield mock_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_connections_client():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.ConnectionsClient"
|
||||
) as mock_client:
|
||||
yield mock_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openapi_toolset():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.OpenAPIToolset"
|
||||
) as mock_toolset:
|
||||
mock_toolset_instance = mock.MagicMock()
|
||||
mock_rest_api_tool = mock.MagicMock(spec=rest_api_tool.RestApiTool)
|
||||
mock_rest_api_tool.name = "Test Tool"
|
||||
|
||||
# Create an async mock for the get_tools method
|
||||
async def mock_get_tools(context: ReadonlyContext = None):
|
||||
return [mock_rest_api_tool]
|
||||
|
||||
# Assign the async mock function to get_tools
|
||||
mock_toolset_instance.get_tools = mock_get_tools
|
||||
|
||||
mock_toolset.return_value = mock_toolset_instance
|
||||
yield mock_toolset
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openapi_toolset_with_multiple_tools_and_no_tools():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.OpenAPIToolset"
|
||||
) as mock_toolset:
|
||||
mock_toolset_instance = mock.MagicMock()
|
||||
mock_rest_api_tool = mock.MagicMock(spec=rest_api_tool.RestApiTool)
|
||||
mock_rest_api_tool.name = "Test Tool"
|
||||
mock_rest_api_tool_2 = mock.MagicMock(spec=rest_api_tool.RestApiTool)
|
||||
mock_rest_api_tool_2.name = "Test Tool 2"
|
||||
|
||||
# Create an async mock for the get_tools method
|
||||
async def mock_get_tools(context: ReadonlyContext = None):
|
||||
return [mock_rest_api_tool, mock_rest_api_tool_2]
|
||||
|
||||
mock_toolset_instance.get_tools = mock_get_tools
|
||||
mock_toolset.return_value = mock_toolset_instance
|
||||
yield mock_toolset
|
||||
|
||||
|
||||
def get_mocked_parsed_operation(operation_id, attributes):
|
||||
mock_openapi_spec_parser_instance = mock.MagicMock()
|
||||
mock_parsed_operation = mock.MagicMock(spec=ParsedOperation)
|
||||
mock_parsed_operation.name = "list_issues"
|
||||
mock_parsed_operation.description = "list_issues_description"
|
||||
mock_parsed_operation.endpoint = OperationEndpoint(
|
||||
base_url="http://localhost:8080",
|
||||
path="/v1/issues",
|
||||
method="GET",
|
||||
)
|
||||
mock_parsed_operation.auth_scheme = None
|
||||
mock_parsed_operation.auth_credential = None
|
||||
mock_parsed_operation.additional_context = {}
|
||||
mock_parsed_operation.parameters = []
|
||||
mock_operation = mock.MagicMock(spec=Operation)
|
||||
mock_operation.operationId = operation_id
|
||||
mock_operation.description = "list_issues_description"
|
||||
mock_operation.parameters = []
|
||||
mock_operation.requestBody = None
|
||||
mock_operation.responses = {}
|
||||
mock_operation.callbacks = {}
|
||||
for key, value in attributes.items():
|
||||
setattr(mock_operation, key, value)
|
||||
mock_parsed_operation.operation = mock_operation
|
||||
mock_openapi_spec_parser_instance.parse.return_value = [mock_parsed_operation]
|
||||
return mock_openapi_spec_parser_instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openapi_entity_spec_parser():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.OpenApiSpecParser"
|
||||
) as mock_spec_parser:
|
||||
mock_openapi_spec_parser_instance = get_mocked_parsed_operation(
|
||||
"list_issues", {"x-entity": "Issues", "x-operation": "LIST_ENTITIES"}
|
||||
)
|
||||
mock_spec_parser.return_value = mock_openapi_spec_parser_instance
|
||||
yield mock_spec_parser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_openapi_action_spec_parser():
|
||||
with mock.patch(
|
||||
"google.adk.tools.application_integration_tool.application_integration_toolset.OpenApiSpecParser"
|
||||
) as mock_spec_parser:
|
||||
mock_openapi_action_spec_parser_instance = get_mocked_parsed_operation(
|
||||
"list_issues_operation",
|
||||
{"x-action": "CustomAction", "x-operation": "EXECUTE_ACTION"},
|
||||
)
|
||||
mock_spec_parser.return_value = mock_openapi_action_spec_parser_instance
|
||||
yield mock_spec_parser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project():
|
||||
return "test-project"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def location():
|
||||
return "us-central1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_spec():
|
||||
return {"openapi": "3.0.0", "info": {"title": "Integration API"}}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_spec():
|
||||
return {"openapi": "3.0.0", "info": {"title": "Connection API"}}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_details():
|
||||
return {
|
||||
"serviceName": "test-service",
|
||||
"host": "test.host",
|
||||
"name": "test-connection",
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_details_auth_override_enabled():
|
||||
return {
|
||||
"serviceName": "test-service",
|
||||
"host": "test.host",
|
||||
"name": "test-connection",
|
||||
"authOverrideEnabled": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_with_integration_and_trigger(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_toolset,
|
||||
):
|
||||
integration_name = "test-integration"
|
||||
triggers = ["test-trigger"]
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project, location, integration=integration_name, triggers=triggers
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
integration_name,
|
||||
triggers,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
|
||||
mock_connections_client.assert_not_called()
|
||||
mock_openapi_toolset.assert_called_once()
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "Test Tool"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_with_integration_and_list_of_triggers(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_toolset_with_multiple_tools_and_no_tools,
|
||||
):
|
||||
integration_name = "test-integration"
|
||||
triggers = ["test-trigger1", "test-trigger2"]
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project, location, integration=integration_name, triggers=triggers
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
integration_name,
|
||||
triggers,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
|
||||
mock_connections_client.assert_not_called()
|
||||
mock_openapi_toolset_with_multiple_tools_and_no_tools.assert_called_once()
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "Test Tool"
|
||||
assert tools[1].name == "Test Tool 2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_with_integration_and_empty_trigger_list(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_toolset_with_multiple_tools_and_no_tools,
|
||||
):
|
||||
integration_name = "test-integration"
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project, location, integration=integration_name
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project, location, None, integration_name, None, None, None, None, None
|
||||
)
|
||||
mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
|
||||
mock_connections_client.assert_not_called()
|
||||
mock_openapi_toolset_with_multiple_tools_and_no_tools.assert_called_once()
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "Test Tool"
|
||||
assert tools[1].name == "Test Tool 2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_with_connection_and_entity_operations(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_entity_spec_parser,
|
||||
connection_details,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
entity_operations_list = ["list", "get"]
|
||||
tool_name = "My Connection Tool"
|
||||
tool_instructions = "Use this tool to manage entities."
|
||||
mock_connections_client.return_value.get_connection_details.return_value = (
|
||||
connection_details
|
||||
)
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
entity_operations=entity_operations_list,
|
||||
tool_name_prefix=tool_name,
|
||||
tool_instructions=tool_instructions,
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
connection_name,
|
||||
entity_operations_list,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_openapi_entity_spec_parser.return_value.parse.assert_called_once()
|
||||
mock_connections_client.return_value.get_connection_details.assert_called_once()
|
||||
mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
|
||||
tool_name,
|
||||
tool_instructions,
|
||||
)
|
||||
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "list_issues"
|
||||
assert isinstance(tools[0], IntegrationConnectorTool)
|
||||
assert tools[0]._entity == "Issues"
|
||||
assert tools[0]._operation == "LIST_ENTITIES"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_initialization_with_connection_and_actions(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_action_spec_parser,
|
||||
connection_details,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
actions_list = ["create", "delete"]
|
||||
tool_name = "My Actions Tool"
|
||||
tool_instructions = "Perform actions using this tool."
|
||||
mock_connections_client.return_value.get_connection_details.return_value = (
|
||||
connection_details
|
||||
)
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
actions=actions_list,
|
||||
tool_name_prefix=tool_name,
|
||||
tool_instructions=tool_instructions,
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
connection_name,
|
||||
None,
|
||||
actions_list,
|
||||
None,
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_connections_client.return_value.get_connection_details.assert_called_once()
|
||||
mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
|
||||
tool_name, tool_instructions
|
||||
)
|
||||
mock_openapi_action_spec_parser.return_value.parse.assert_called_once()
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "list_issues_operation"
|
||||
assert isinstance(tools[0], IntegrationConnectorTool)
|
||||
assert tools[0]._action == "CustomAction"
|
||||
assert tools[0]._operation == "EXECUTE_ACTION"
|
||||
|
||||
|
||||
def test_initialization_without_required_params(project, location):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Invalid request, Either integration or \\(connection and"
|
||||
" \\(entity_operations or actions\\)\\) should be provided."
|
||||
),
|
||||
):
|
||||
ApplicationIntegrationToolset(project, location)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Invalid request, Either integration or \\(connection and"
|
||||
" \\(entity_operations or actions\\)\\) should be provided."
|
||||
),
|
||||
):
|
||||
ApplicationIntegrationToolset(project, location, triggers=["test"])
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Invalid request, Either integration or \\(connection and"
|
||||
" \\(entity_operations or actions\\)\\) should be provided."
|
||||
),
|
||||
):
|
||||
ApplicationIntegrationToolset(project, location, connection="test")
|
||||
|
||||
|
||||
def test_initialization_with_service_account_credentials(
|
||||
project, location, mock_integration_client, mock_openapi_toolset
|
||||
):
|
||||
service_account_json = json.dumps({
|
||||
"type": "service_account",
|
||||
"project_id": "dummy",
|
||||
"private_key_id": "dummy",
|
||||
"private_key": "dummy",
|
||||
"client_email": "test@example.com",
|
||||
"client_id": "131331543646416",
|
||||
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
|
||||
"token_uri": "https://oauth2.googleapis.com/token",
|
||||
"auth_provider_x509_cert_url": (
|
||||
"https://www.googleapis.com/oauth2/v1/certs"
|
||||
),
|
||||
"client_x509_cert_url": (
|
||||
"http://www.googleapis.com/robot/v1/metadata/x509/dummy%40dummy.com"
|
||||
),
|
||||
"universe_domain": "googleapis.com",
|
||||
})
|
||||
integration_name = "test-integration"
|
||||
triggers = ["test-trigger"]
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
integration=integration_name,
|
||||
triggers=triggers,
|
||||
service_account_json=service_account_json,
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
integration_name,
|
||||
triggers,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
service_account_json,
|
||||
)
|
||||
mock_openapi_toolset.assert_called_once()
|
||||
_, kwargs = mock_openapi_toolset.call_args
|
||||
assert isinstance(kwargs["auth_credential"], AuthCredential)
|
||||
assert (
|
||||
kwargs[
|
||||
"auth_credential"
|
||||
].service_account.service_account_credential.client_email
|
||||
== "test@example.com"
|
||||
)
|
||||
|
||||
|
||||
def test_initialization_without_explicit_service_account_credentials(
|
||||
project, location, mock_integration_client, mock_openapi_toolset
|
||||
):
|
||||
integration_name = "test-integration"
|
||||
triggers = "test-trigger"
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project, location, integration=integration_name, triggers=triggers
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
integration_name,
|
||||
triggers,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
mock_openapi_toolset.assert_called_once()
|
||||
_, kwargs = mock_openapi_toolset.call_args
|
||||
assert isinstance(kwargs["auth_credential"], AuthCredential)
|
||||
assert kwargs["auth_credential"].service_account.use_default_credential
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tools(
|
||||
project, location, mock_integration_client, mock_openapi_toolset
|
||||
):
|
||||
integration_name = "test-integration"
|
||||
triggers = ["test-trigger"]
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project, location, integration=integration_name, triggers=triggers
|
||||
)
|
||||
tools = await toolset.get_tools()
|
||||
assert len(tools) == 1
|
||||
assert isinstance(tools[0], rest_api_tool.RestApiTool)
|
||||
assert tools[0].name == "Test Tool"
|
||||
|
||||
|
||||
def test_initialization_with_connection_details(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_toolset,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
entity_operations_list = ["list"]
|
||||
tool_name = "My Connection Tool"
|
||||
tool_instructions = "Use this tool."
|
||||
mock_connections_client.return_value.get_connection_details.return_value = {
|
||||
"serviceName": "custom-service",
|
||||
"host": "custom.host",
|
||||
}
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
entity_operations=entity_operations_list,
|
||||
tool_name_prefix=tool_name,
|
||||
tool_instructions=tool_instructions,
|
||||
)
|
||||
mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
|
||||
tool_name, tool_instructions
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_with_connection_and_custom_auth(
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_action_spec_parser,
|
||||
connection_details_auth_override_enabled,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
actions_list = ["create", "delete"]
|
||||
tool_name = "My Actions Tool"
|
||||
tool_instructions = "Perform actions using this tool."
|
||||
mock_connections_client.return_value.get_connection_details.return_value = (
|
||||
connection_details_auth_override_enabled
|
||||
)
|
||||
|
||||
oauth2_data_google_cloud = {
|
||||
"type": "oauth2",
|
||||
"flows": {
|
||||
"authorizationCode": {
|
||||
"authorizationUrl": "https://test-url/o/oauth2/auth",
|
||||
"tokenUrl": "https://test-url/token",
|
||||
"scopes": {
|
||||
"https://test-url/auth/test-scope": "test scope",
|
||||
"https://www.test-url.com/auth/test-scope2": "test scope 2",
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
oauth2_scheme = dict_to_auth_scheme(oauth2_data_google_cloud)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test-client-id",
|
||||
client_secret="test-client-secret",
|
||||
),
|
||||
)
|
||||
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
actions=actions_list,
|
||||
tool_name_prefix=tool_name,
|
||||
tool_instructions=tool_instructions,
|
||||
auth_scheme=oauth2_scheme,
|
||||
auth_credential=auth_credential,
|
||||
credential_key="test-key",
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
connection_name,
|
||||
None,
|
||||
actions_list,
|
||||
None,
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_connections_client.return_value.get_connection_details.assert_called_once()
|
||||
mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
|
||||
tool_name, tool_instructions
|
||||
)
|
||||
mock_openapi_action_spec_parser.return_value.parse.assert_called_once()
|
||||
assert len(await toolset.get_tools()) == 1
|
||||
assert (await toolset.get_tools())[0].name == "list_issues_operation"
|
||||
assert isinstance((await toolset.get_tools())[0], IntegrationConnectorTool)
|
||||
assert (await toolset.get_tools())[0]._action == "CustomAction"
|
||||
assert (await toolset.get_tools())[0]._operation == "EXECUTE_ACTION"
|
||||
assert (await toolset.get_tools())[0]._auth_scheme == oauth2_scheme
|
||||
assert (await toolset.get_tools())[0]._auth_credential == auth_credential
|
||||
assert (await toolset.get_tools())[0]._credential_key == "test-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_with_connection_with_auth_override_disabled_and_custom_auth(
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_action_spec_parser,
|
||||
connection_details,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
actions_list = ["create", "delete"]
|
||||
tool_name = "My Actions Tool"
|
||||
tool_instructions = "Perform actions using this tool."
|
||||
mock_connections_client.return_value.get_connection_details.return_value = (
|
||||
connection_details
|
||||
)
|
||||
|
||||
oauth2_data_google_cloud = {
|
||||
"type": "oauth2",
|
||||
"flows": {
|
||||
"authorizationCode": {
|
||||
"authorizationUrl": "https://test-url/o/oauth2/auth",
|
||||
"tokenUrl": "https://test-url/token",
|
||||
"scopes": {
|
||||
"https://test-url/auth/test-scope": "test scope",
|
||||
"https://www.test-url.com/auth/test-scope2": "test scope 2",
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
oauth2_scheme = dict_to_auth_scheme(oauth2_data_google_cloud)
|
||||
|
||||
auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test-client-id",
|
||||
client_secret="test-client-secret",
|
||||
),
|
||||
)
|
||||
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
actions=actions_list,
|
||||
tool_name_prefix=tool_name,
|
||||
tool_instructions=tool_instructions,
|
||||
auth_scheme=oauth2_scheme,
|
||||
auth_credential=auth_credential,
|
||||
)
|
||||
mock_integration_client.assert_called_once_with(
|
||||
project,
|
||||
location,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
connection_name,
|
||||
None,
|
||||
actions_list,
|
||||
None,
|
||||
)
|
||||
mock_connections_client.assert_called_once_with(
|
||||
project, location, connection_name, None
|
||||
)
|
||||
mock_connections_client.return_value.get_connection_details.assert_called_once()
|
||||
mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
|
||||
tool_name, tool_instructions
|
||||
)
|
||||
mock_openapi_action_spec_parser.return_value.parse.assert_called_once()
|
||||
assert len(await toolset.get_tools()) == 1
|
||||
assert (await toolset.get_tools())[0].name == "list_issues_operation"
|
||||
assert isinstance((await toolset.get_tools())[0], IntegrationConnectorTool)
|
||||
assert (await toolset.get_tools())[0]._action == "CustomAction"
|
||||
assert (await toolset.get_tools())[0]._operation == "EXECUTE_ACTION"
|
||||
assert not (await toolset.get_tools())[0]._auth_scheme
|
||||
assert not (await toolset.get_tools())[0]._auth_credential
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tools_uses_exchanged_auth_credential_when_available(
|
||||
project,
|
||||
location,
|
||||
mock_integration_client,
|
||||
mock_connections_client,
|
||||
mock_openapi_action_spec_parser,
|
||||
connection_details_auth_override_enabled,
|
||||
):
|
||||
connection_name = "test-connection"
|
||||
actions_list = ["create"]
|
||||
mock_connections_client.return_value.get_connection_details.return_value = (
|
||||
connection_details_auth_override_enabled
|
||||
)
|
||||
|
||||
oauth2_data_google_cloud = {
|
||||
"type": "oauth2",
|
||||
"flows": {
|
||||
"authorizationCode": {
|
||||
"authorizationUrl": "https://test-url/o/oauth2/auth",
|
||||
"tokenUrl": "https://test-url/token",
|
||||
"scopes": {
|
||||
"https://test-url/auth/test-scope": "test scope",
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
oauth2_scheme = dict_to_auth_scheme(oauth2_data_google_cloud)
|
||||
raw_auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test-client-id",
|
||||
client_secret="test-client-secret",
|
||||
),
|
||||
)
|
||||
|
||||
toolset = ApplicationIntegrationToolset(
|
||||
project,
|
||||
location,
|
||||
connection=connection_name,
|
||||
actions=actions_list,
|
||||
auth_scheme=oauth2_scheme,
|
||||
auth_credential=raw_auth_credential,
|
||||
)
|
||||
|
||||
exchanged_auth_credential = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.OAUTH2,
|
||||
oauth2=OAuth2Auth(
|
||||
client_id="test-client-id",
|
||||
client_secret="test-client-secret",
|
||||
access_token="exchanged-access-token",
|
||||
),
|
||||
)
|
||||
toolset._auth_config.exchanged_auth_credential = exchanged_auth_credential
|
||||
|
||||
original_tool = toolset._tools[0]
|
||||
tools = await toolset.get_tools()
|
||||
|
||||
assert len(tools) == 1
|
||||
assert tools[0] is not original_tool
|
||||
assert tools[0]._auth_credential == exchanged_auth_credential
|
||||
assert original_tool._auth_credential == raw_auth_credential
|
||||
+308
@@ -0,0 +1,308 @@
|
||||
# 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.
|
||||
|
||||
from unittest import mock
|
||||
|
||||
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.features import FeatureName
|
||||
from google.adk.features._feature_registry import temporary_feature_override
|
||||
from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult
|
||||
from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler
|
||||
from google.genai.types import FunctionDeclaration
|
||||
from google.genai.types import Schema
|
||||
from google.genai.types import Type
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_rest_api_tool():
|
||||
"""Fixture for a mocked RestApiTool."""
|
||||
mock_tool = mock.MagicMock(spec=RestApiTool)
|
||||
mock_tool.name = "mock_rest_tool"
|
||||
mock_tool.description = "Mock REST tool description."
|
||||
# Mock the internal parser needed for _get_declaration
|
||||
mock_parser = mock.MagicMock()
|
||||
mock_parser.get_json_schema.return_value = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "string", "description": "User ID"},
|
||||
"connection_name": {"type": "string"},
|
||||
"host": {"type": "string"},
|
||||
"service_name": {"type": "string"},
|
||||
"entity": {"type": "string"},
|
||||
"operation": {"type": "string"},
|
||||
"action": {"type": "string"},
|
||||
"page_size": {"type": "integer"},
|
||||
"filter": {"type": "string"},
|
||||
},
|
||||
"required": ["user_id", "page_size", "filter", "connection_name"],
|
||||
}
|
||||
mock_tool._operation_parser = mock_parser
|
||||
mock_tool.call = mock.AsyncMock(
|
||||
return_value={"status": "success", "data": "mock_data"}
|
||||
)
|
||||
return mock_tool
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_tool(mock_rest_api_tool):
|
||||
"""Fixture for an IntegrationConnectorTool instance."""
|
||||
return IntegrationConnectorTool(
|
||||
name="test_integration_tool",
|
||||
description="Test integration tool description.",
|
||||
connection_name="test-conn",
|
||||
connection_host="test.example.com",
|
||||
connection_service_name="test-service",
|
||||
entity="TestEntity",
|
||||
operation="LIST",
|
||||
action="TestAction",
|
||||
rest_api_tool=mock_rest_api_tool,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_tool_with_auth(mock_rest_api_tool):
|
||||
"""Fixture for an IntegrationConnectorTool instance."""
|
||||
return IntegrationConnectorTool(
|
||||
name="test_integration_tool",
|
||||
description="Test integration tool description.",
|
||||
connection_name="test-conn",
|
||||
connection_host="test.example.com",
|
||||
connection_service_name="test-service",
|
||||
entity="TestEntity",
|
||||
operation="LIST",
|
||||
action="TestAction",
|
||||
rest_api_tool=mock_rest_api_tool,
|
||||
auth_scheme=None,
|
||||
auth_credential=AuthCredential(
|
||||
auth_type=AuthCredentialTypes.HTTP,
|
||||
http=HttpAuth(
|
||||
scheme="bearer",
|
||||
credentials=HttpCredentials(token="mocked_token"),
|
||||
),
|
||||
),
|
||||
credential_key="test-key",
|
||||
)
|
||||
|
||||
|
||||
class TestIntegrationConnectorToolLegacy:
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def disable_feature_flag(self):
|
||||
"""Disable the JSON_SCHEMA_FOR_FUNC_DECL feature flag for legacy tests."""
|
||||
with temporary_feature_override(
|
||||
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, False
|
||||
):
|
||||
yield
|
||||
|
||||
def test_get_declaration(self, integration_tool):
|
||||
"""Tests the generation of the function declaration."""
|
||||
declaration = integration_tool._get_declaration()
|
||||
|
||||
assert isinstance(declaration, FunctionDeclaration)
|
||||
assert declaration.name == "test_integration_tool"
|
||||
assert declaration.description == "Test integration tool description."
|
||||
|
||||
# Check parameters schema
|
||||
params = declaration.parameters
|
||||
assert isinstance(params, Schema)
|
||||
print(f"params: {params}")
|
||||
assert params.type == Type.OBJECT
|
||||
|
||||
# Check properties (excluded fields should not be present)
|
||||
assert "user_id" in params.properties
|
||||
assert "connection_name" not in params.properties
|
||||
assert "host" not in params.properties
|
||||
assert "service_name" not in params.properties
|
||||
assert "entity" not in params.properties
|
||||
assert "operation" not in params.properties
|
||||
assert "action" not in params.properties
|
||||
assert "page_size" in params.properties
|
||||
assert "filter" in params.properties
|
||||
|
||||
# Check required fields (optional and excluded fields should not be required)
|
||||
assert "user_id" in params.required
|
||||
assert "page_size" not in params.required
|
||||
assert "filter" not in params.required
|
||||
assert "connection_name" not in params.required
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async(integration_tool, mock_rest_api_tool):
|
||||
"""Tests the async execution delegates correctly to the RestApiTool."""
|
||||
input_args = {"user_id": "user123", "page_size": 10}
|
||||
expected_call_args = {
|
||||
"user_id": "user123",
|
||||
"page_size": 10,
|
||||
"connection_name": "test-conn",
|
||||
"host": "test.example.com",
|
||||
"service_name": "test-service",
|
||||
"entity": "TestEntity",
|
||||
"operation": "LIST",
|
||||
"action": "TestAction",
|
||||
}
|
||||
|
||||
result = await integration_tool.run_async(args=input_args, tool_context=None)
|
||||
|
||||
# Assert the underlying rest_api_tool.call was called correctly
|
||||
mock_rest_api_tool.call.assert_called_once_with(
|
||||
args=expected_call_args, tool_context=None
|
||||
)
|
||||
|
||||
# Assert the result is what the mocked call returned
|
||||
assert result == {"status": "success", "data": "mock_data"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_auth_async_none_token(
|
||||
integration_tool_with_auth, mock_rest_api_tool
|
||||
):
|
||||
"""Tests run_async when auth credential token is None."""
|
||||
input_args = {
|
||||
"user_id": "user456",
|
||||
"filter": "some_filter",
|
||||
"sortByColumns": ["a", "b"],
|
||||
}
|
||||
expected_call_args = {
|
||||
"user_id": "user456",
|
||||
"filter": "some_filter",
|
||||
"dynamic_auth_config": {"oauth2_auth_code_flow.access_token": {}},
|
||||
"connection_name": "test-conn",
|
||||
"service_name": "test-service",
|
||||
"host": "test.example.com",
|
||||
"entity": "TestEntity",
|
||||
"operation": "LIST",
|
||||
"action": "TestAction",
|
||||
"sortByColumns": ["a", "b"],
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
ToolAuthHandler, "from_tool_context", autospec=True
|
||||
) as mock_from_tool_context:
|
||||
mock_tool_auth_handler_instance = mock.MagicMock()
|
||||
# Simulate an AuthCredential that would cause _prepare_dynamic_euc to return None
|
||||
mock_auth_credential_without_token = AuthCredential(
|
||||
auth_type=AuthCredentialTypes.HTTP,
|
||||
http=HttpAuth(
|
||||
scheme="bearer",
|
||||
credentials=HttpCredentials(token=None), # Token is None
|
||||
),
|
||||
)
|
||||
mock_tool_auth_handler_instance.prepare_auth_credentials = mock.AsyncMock(
|
||||
return_value=(
|
||||
AuthPreparationResult(
|
||||
state="done", auth_credential=mock_auth_credential_without_token
|
||||
)
|
||||
)
|
||||
)
|
||||
mock_from_tool_context.return_value = mock_tool_auth_handler_instance
|
||||
|
||||
result = await integration_tool_with_auth.run_async(
|
||||
args=input_args, tool_context={}
|
||||
)
|
||||
mock_from_tool_context.assert_called_once_with(
|
||||
{},
|
||||
None,
|
||||
integration_tool_with_auth._auth_credential,
|
||||
credential_key="test-key",
|
||||
)
|
||||
|
||||
mock_rest_api_tool.call.assert_called_once_with(
|
||||
args=expected_call_args, tool_context={}
|
||||
)
|
||||
assert result == {"status": "success", "data": "mock_data"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_auth_async(
|
||||
integration_tool_with_auth, mock_rest_api_tool
|
||||
):
|
||||
"""Tests the async execution with auth delegates correctly to the RestApiTool."""
|
||||
input_args = {"user_id": "user123", "page_size": 10}
|
||||
expected_call_args = {
|
||||
"user_id": "user123",
|
||||
"page_size": 10,
|
||||
"dynamic_auth_config": {
|
||||
"oauth2_auth_code_flow.access_token": "mocked_token"
|
||||
},
|
||||
"connection_name": "test-conn",
|
||||
"service_name": "test-service",
|
||||
"host": "test.example.com",
|
||||
"entity": "TestEntity",
|
||||
"operation": "LIST",
|
||||
"action": "TestAction",
|
||||
}
|
||||
|
||||
with mock.patch.object(
|
||||
ToolAuthHandler, "from_tool_context", autospec=True
|
||||
) as mock_from_tool_context:
|
||||
mock_tool_auth_handler_instance = mock.MagicMock()
|
||||
|
||||
mock_tool_auth_handler_instance.prepare_auth_credentials = mock.AsyncMock(
|
||||
return_value=AuthPreparationResult(
|
||||
state="done",
|
||||
auth_credential=AuthCredential(
|
||||
auth_type=AuthCredentialTypes.HTTP,
|
||||
http=HttpAuth(
|
||||
scheme="bearer",
|
||||
credentials=HttpCredentials(token="mocked_token"),
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
mock_from_tool_context.return_value = mock_tool_auth_handler_instance
|
||||
result = await integration_tool_with_auth.run_async(
|
||||
args=input_args, tool_context={}
|
||||
)
|
||||
mock_from_tool_context.assert_called_once_with(
|
||||
{},
|
||||
None,
|
||||
integration_tool_with_auth._auth_credential,
|
||||
credential_key="test-key",
|
||||
)
|
||||
mock_rest_api_tool.call.assert_called_once_with(
|
||||
args=expected_call_args, tool_context={}
|
||||
)
|
||||
assert result == {"status": "success", "data": "mock_data"}
|
||||
|
||||
|
||||
class TestIntegrationConnectorToolWithJsonSchema:
|
||||
|
||||
def test_get_declaration_with_json_schema_feature_enabled(
|
||||
self, integration_tool
|
||||
):
|
||||
"""Tests the generation of the function declaration with JSON schema feature enabled."""
|
||||
with temporary_feature_override(
|
||||
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
|
||||
):
|
||||
declaration = integration_tool._get_declaration()
|
||||
|
||||
assert isinstance(declaration, FunctionDeclaration)
|
||||
assert declaration.name == "test_integration_tool"
|
||||
assert declaration.description == "Test integration tool description."
|
||||
assert declaration.parameters is None
|
||||
assert declaration.parameters_json_schema == {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "string", "description": "User ID"},
|
||||
"page_size": {"type": "integer"},
|
||||
"filter": {"type": "string"},
|
||||
},
|
||||
"required": ["user_id"],
|
||||
}
|
||||
Reference in New Issue
Block a user