google-adk 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

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

@@ -1,600 +0,0 @@
-# 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()