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

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

@@ -1,630 +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
-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 trigger_name():
-  return "test-trigger"
-
-
-@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, trigger_name, connection_name
-  ):
-    client = IntegrationClient(
-        project=project,
-        location=location,
-        integration=integration_name,
-        trigger=trigger_name,
-        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.trigger == trigger_name
-    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,
-      trigger_name,
-      mock_credentials,
-      mock_connections_client,
-  ):
-    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(
-        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,
-          trigger=trigger_name,
-          connection=None,
-          entity_operations=None,
-          actions=None,
-          service_account_json=None,
-      )
-      spec = client.get_openapi_spec_for_integration()
-      assert spec == expected_spec
-      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}",
-          },
-          json={
-              "apiTriggerResources": [{
-                  "integrationResource": integration_name,
-                  "triggerId": [trigger_name],
-              }],
-              "fileFormat": "JSON",
-          },
-      )
-
-  def test_get_openapi_spec_for_integration_credential_error(
-      self,
-      project,
-      location,
-      integration_name,
-      trigger_name,
-      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,
-          trigger=trigger_name,
-          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,
-      trigger_name,
-      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,
-          trigger=trigger_name,
-          connection=None,
-          entity_operations=None,
-          actions=None,
-          service_account_json=None,
-      )
-      with pytest.raises(
-          ValueError,
-          match=(
-              "Invalid request. Please check the provided values of"
-              f" project\\({project}\\), location\\({location}\\),"
-              f" integration\\({integration_name}\\) and"
-              f" trigger\\({trigger_name}\\)."
-          ),
-      ):
-        client.get_openapi_spec_for_integration()
-
-  def test_get_openapi_spec_for_integration_other_request_error(
-      self,
-      project,
-      location,
-      integration_name,
-      trigger_name,
-      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,
-          trigger=trigger_name,
-          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,
-      trigger_name,
-      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,
-          trigger=trigger_name,
-          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,
-        trigger=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,
-        trigger=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,
-        trigger=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,
-        trigger=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, trigger_name, 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,
-          trigger=trigger_name,
-          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,
-      trigger_name,
-      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"),
-    ), mock.patch.object(mock_credentials, "refresh", return_value=None):
-      client = IntegrationClient(
-          project=project,
-          location=location,
-          integration=integration_name,
-          trigger=trigger_name,
-          connection=connection_name,
-          entity_operations=None,
-          actions=None,
-          service_account_json=None,
-      )
-      token = client._get_access_token()
-      assert token == "test_token"
-
-  def test_get_access_token_no_valid_credentials(
-      self, project, location, integration_name, trigger_name, 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,
-          trigger=trigger_name,
-          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_uses_cached_token(
-      self,
-      project,
-      location,
-      integration_name,
-      trigger_name,
-      connection_name,
-      mock_credentials,
-  ):
-    mock_credentials.token = "cached_token"
-    mock_credentials.expired = False
-    client = IntegrationClient(
-        project=project,
-        location=location,
-        integration=integration_name,
-        trigger=trigger_name,
-        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,
-      trigger_name,
-      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,
-          trigger=trigger_name,
-          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()