Skip to content

Commit a9007fb

Browse files
its-amannits-amann
committed
Merge remote-tracking branch 'origin/main' into fix-thought-signature-pruning
2 parents cb0cdb9 + 37ca6fb commit a9007fb

4 files changed

Lines changed: 149 additions & 14 deletions

File tree

‎src/google/adk/tools/application_integration_tool/clients/connections_client.py‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,15 @@
2828
from google.oauth2 import service_account
2929
import requests
3030

31+
from ....utils import _mtls_utils
32+
33+
_DEFAULT_CONNECTORS_ENDPOINT_TEMPLATE = "connectors.googleapis.com"
34+
_DEFAULT_MTLS_CONNECTORS_ENDPOINT_TEMPLATE = "connectors.mtls.googleapis.com"
35+
_DEFAULT_INTEGRATIONS_ENDPOINT_TEMPLATE = "integrations.googleapis.com"
36+
_DEFAULT_MTLS_INTEGRATIONS_ENDPOINT_TEMPLATE = (
37+
"integrations.mtls.googleapis.com"
38+
)
39+
3140

3241
class ConnectionsClient:
3342
"""Utility class for interacting with Google Cloud Connectors API."""
@@ -52,7 +61,11 @@ def __init__(
5261
self.project = project
5362
self.location = location
5463
self.connection = connection
55-
self.connector_url = "https://connectors.googleapis.com"
64+
self.connector_url = "https://" + _mtls_utils.get_api_endpoint(
65+
location,
66+
_DEFAULT_CONNECTORS_ENDPOINT_TEMPLATE,
67+
_DEFAULT_MTLS_CONNECTORS_ENDPOINT_TEMPLATE,
68+
)
5669
self.service_account_json = service_account_json
5770
self.credential_cache = None
5871

@@ -168,7 +181,16 @@ def get_connector_base_spec() -> Dict[str, Any]:
168181
"description": "This tool can execute a query on connection",
169182
"version": "4",
170183
},
171-
"servers": [{"url": "https://integrations.googleapis.com"}],
184+
"servers": [{
185+
"url": (
186+
"https://"
187+
+ _mtls_utils.get_api_endpoint(
188+
"",
189+
_DEFAULT_INTEGRATIONS_ENDPOINT_TEMPLATE,
190+
_DEFAULT_MTLS_INTEGRATIONS_ENDPOINT_TEMPLATE,
191+
)
192+
)
193+
}],
172194
"security": [
173195
{"google_auth": ["https://www.googleapis.com/auth/cloud-platform"]}
174196
],

‎src/google/adk/tools/application_integration_tool/clients/integration_client.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,15 @@
2626
from google.oauth2 import service_account
2727
import requests
2828

29+
from ....utils import _mtls_utils
30+
31+
_DEFAULT_INTEGRATIONS_REGIONAL_ENDPOINT_TEMPLATE = (
32+
"{location}-integrations.googleapis.com"
33+
)
34+
_DEFAULT_MTLS_INTEGRATIONS_REGIONAL_ENDPOINT_TEMPLATE = (
35+
"{location}-integrations.mtls.googleapis.com"
36+
)
37+
2938

3039
class IntegrationClient:
3140
"""A client for interacting with Google Cloud Application Integration.
@@ -88,7 +97,12 @@ def get_openapi_spec_for_integration(self):
8897
Exception: For any other unexpected errors.
8998
"""
9099
try:
91-
url = f"https://{self.location}-integrations.googleapis.com/v1/projects/{self.project}/locations/{self.location}:generateOpenApiSpec"
100+
endpoint = _mtls_utils.get_api_endpoint(
101+
self.location,
102+
_DEFAULT_INTEGRATIONS_REGIONAL_ENDPOINT_TEMPLATE,
103+
_DEFAULT_MTLS_INTEGRATIONS_REGIONAL_ENDPOINT_TEMPLATE,
104+
)
105+
url = f"https://{endpoint}/v1/projects/{self.project}/locations/{self.location}:generateOpenApiSpec"
92106
headers = {
93107
"Content-Type": "application/json",
94108
"Authorization": f"Bearer {self._get_access_token()}",

‎tests/unittests/tools/application_integration_tool/clients/test_connections_client.py‎

Lines changed: 45 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -54,15 +54,37 @@ class TestConnectionsClient:
5454

5555
def test_initialization(self, project, location, connection_name):
5656
credentials = {"email": "test@example.com"}
57-
client = ConnectionsClient(
58-
project, location, connection_name, json.dumps(credentials)
59-
)
60-
assert client.project == project
61-
assert client.location == location
62-
assert client.connection == connection_name
63-
assert client.connector_url == "https://connectors.googleapis.com"
64-
assert client.service_account_json == json.dumps(credentials)
65-
assert client.credential_cache is None
57+
with mock.patch(
58+
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
59+
return_value="connectors.googleapis.com",
60+
) as mock_get_api_endpoint:
61+
client = ConnectionsClient(
62+
project, location, connection_name, json.dumps(credentials)
63+
)
64+
assert client.project == project
65+
assert client.location == location
66+
assert client.connection == connection_name
67+
assert client.connector_url == "https://connectors.googleapis.com"
68+
assert client.service_account_json == json.dumps(credentials)
69+
assert client.credential_cache is None
70+
mock_get_api_endpoint.assert_called_once_with(
71+
location,
72+
"connectors.googleapis.com",
73+
"connectors.mtls.googleapis.com",
74+
)
75+
76+
def test_initialization_mtls_endpoint(
77+
self, project, location, connection_name
78+
):
79+
credentials = {"email": "test@example.com"}
80+
with mock.patch(
81+
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
82+
return_value="connectors.mtls.googleapis.com",
83+
):
84+
client = ConnectionsClient(
85+
project, location, connection_name, json.dumps(credentials)
86+
)
87+
assert client.connector_url == "https://connectors.mtls.googleapis.com"
6688

6789
def test_execute_api_call_success(
6890
self, project, location, connection_name, mock_credentials
@@ -261,7 +283,13 @@ def test_get_entity_schema_and_operations_success(
261283
self, project, location, connection_name, mock_credentials
262284
):
263285
credentials = {"email": "test@example.com"}
264-
client = ConnectionsClient(project, location, connection_name, credentials)
286+
with mock.patch(
287+
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
288+
return_value="connectors.googleapis.com",
289+
):
290+
client = ConnectionsClient(
291+
project, location, connection_name, credentials
292+
)
265293
mock_execute_response_initial = mock.MagicMock()
266294
mock_execute_response_initial.status_code = 200
267295
mock_execute_response_initial.json.return_value = {
@@ -335,7 +363,13 @@ def test_get_action_schema_success(
335363
self, project, location, connection_name, mock_credentials
336364
):
337365
credentials = {"email": "test@example.com"}
338-
client = ConnectionsClient(project, location, connection_name, credentials)
366+
with mock.patch(
367+
"google.adk.tools.application_integration_tool.clients.connections_client._mtls_utils.get_api_endpoint",
368+
return_value="connectors.googleapis.com",
369+
):
370+
client = ConnectionsClient(
371+
project, location, connection_name, credentials
372+
)
339373
mock_execute_response_initial = mock.MagicMock()
340374
mock_execute_response_initial.status_code = 200
341375
mock_execute_response_initial.json.return_value = {

‎tests/unittests/tools/application_integration_tool/clients/test_integration_client.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,10 @@ def test_get_openapi_spec_for_integration_success(
126126
),
127127
mock.patch.object(mock_credentials, "refresh", return_value=None),
128128
mock.patch.object(requests, "post", return_value=mock_response),
129+
mock.patch(
130+
"google.adk.tools.application_integration_tool.clients.integration_client._mtls_utils.get_api_endpoint",
131+
return_value=f"{location}-integrations.googleapis.com",
132+
) as mock_get_api_endpoint,
129133
):
130134
client = IntegrationClient(
131135
project=project,
@@ -139,6 +143,11 @@ def test_get_openapi_spec_for_integration_success(
139143
)
140144
spec = client.get_openapi_spec_for_integration()
141145
assert spec == expected_spec
146+
mock_get_api_endpoint.assert_called_once_with(
147+
location,
148+
"{location}-integrations.googleapis.com",
149+
"{location}-integrations.mtls.googleapis.com",
150+
)
142151
requests.post.assert_called_once_with(
143152
f"https://{location}-integrations.googleapis.com/v1/projects/{project}/locations/{location}:generateOpenApiSpec",
144153
headers={
@@ -155,6 +164,62 @@ def test_get_openapi_spec_for_integration_success(
155164
},
156165
)
157166

167+
def test_get_openapi_spec_for_integration_success_mtls(
168+
self,
169+
project,
170+
location,
171+
integration_name,
172+
triggers,
173+
mock_credentials,
174+
mock_connections_client,
175+
):
176+
mock_credentials.quota_project_id = "quota-project"
177+
mock_credentials.expired = False
178+
expected_spec = {"openapi": "3.0.0", "info": {"title": "Test Integration"}}
179+
mock_response = mock.MagicMock()
180+
mock_response.status_code = 200
181+
mock_response.json.return_value = {"openApiSpec": json.dumps(expected_spec)}
182+
183+
with (
184+
mock.patch.object(
185+
integration_client,
186+
"default_service_credential",
187+
return_value=(mock_credentials, project),
188+
),
189+
mock.patch.object(mock_credentials, "refresh", return_value=None),
190+
mock.patch.object(requests, "post", return_value=mock_response),
191+
mock.patch(
192+
"google.adk.tools.application_integration_tool.clients.integration_client._mtls_utils.get_api_endpoint",
193+
return_value=f"{location}-integrations.mtls.googleapis.com",
194+
),
195+
):
196+
client = IntegrationClient(
197+
project=project,
198+
location=location,
199+
integration=integration_name,
200+
triggers=triggers,
201+
connection=None,
202+
entity_operations=None,
203+
actions=None,
204+
service_account_json=None,
205+
)
206+
client.get_openapi_spec_for_integration()
207+
requests.post.assert_called_once_with(
208+
f"https://{location}-integrations.mtls.googleapis.com/v1/projects/{project}/locations/{location}:generateOpenApiSpec",
209+
headers={
210+
"Content-Type": "application/json",
211+
"Authorization": f"Bearer {mock_credentials.token}",
212+
"x-goog-user-project": "quota-project",
213+
},
214+
json={
215+
"apiTriggerResources": [{
216+
"integrationResource": integration_name,
217+
"triggerId": triggers,
218+
}],
219+
"fileFormat": "JSON",
220+
},
221+
)
222+
158223
def test_get_openapi_spec_for_integration_credential_error(
159224
self,
160225
project,

0 commit comments

Comments
 (0)