google-adk 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

google/adk/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py

@@ -0,0 +1,600 @@
+# Copyright 2025 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"}
+    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
+
+  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 = {
+        "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 == {
+          "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 = {
+        "serviceDirectory": "test_service",
+        "authOverrideEnabled": False,
+    }
+
+    with mock.patch.object(
+        client, "_execute_api_call", return_value=mock_response
+    ):
+      details = client.get_connection_details()
+      assert details == {
+          "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"}
+    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"}
+    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"] == "Create 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"] == "Update 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 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"]
+
+  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"]
+
+  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"]
+
+  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"),
+    ), mock.patch.object(mock_credentials, "refresh", return_value=None):
+      token = client._get_access_token()
+      assert token == "test_token"
+
+  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_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()