Support async agent and model callbacks

PiperOrigin-RevId: 755542756
2025-12-18 11:22:22 -06:00 · 2025-05-06 15:13:39 -07:00
parent f96cdc675c
commit 794a70edcd
25 changed files with 371 additions and 117 deletions
--- a/tests/unittests/tools/init.py
+++ b/tests/unittests/tools/init.py
@@ -11,4 +11,3 @@
 # 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.
-
--- a/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py
+++ b/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py
@@ -74,9 +74,12 @@ class TestConnectionsClient:
    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):
+    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(
@@ -121,9 +124,12 @@ class TestConnectionsClient:
        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 (
+        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"
      ):
@@ -140,9 +146,12 @@ class TestConnectionsClient:
        "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 (
+        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")

@@ -151,10 +160,13 @@ class TestConnectionsClient:
  ):
    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 (
+        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"
@@ -539,10 +551,13 @@ class TestConnectionsClient:
    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):
+    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(
@@ -555,10 +570,13 @@ class TestConnectionsClient:
      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):
+    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"

--- a/tests/unittests/tools/application_integration_tool/clients/test_integration_client.py
+++ b/tests/unittests/tools/application_integration_tool/clients/test_integration_client.py
@@ -114,11 +114,14 @@ class TestIntegrationClient:
    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):
+    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,
@@ -202,11 +205,14 @@ class TestIntegrationClient:
        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):
+    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,
@@ -243,11 +249,14 @@ class TestIntegrationClient:
        "Internal Server Error"
    )

-    with mock.patch.object(
-        IntegrationClient,
-        "_get_access_token",
-        return_value=mock_credentials.token,
-    ), mock.patch("requests.post", return_value=mock_response):
+    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,
@@ -270,12 +279,15 @@ class TestIntegrationClient:
      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")
+    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,
@@ -486,10 +498,13 @@ class TestIntegrationClient:
    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):
+    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,
@@ -518,10 +533,13 @@ class TestIntegrationClient:
      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):
+    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,
@@ -538,12 +556,15 @@ class TestIntegrationClient:
  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,
+    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,
@@ -587,9 +608,12 @@ class TestIntegrationClient:
        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:
+    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()
--- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py
+++ b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py
@@ -236,8 +236,7 @@ def test_initialization_with_connection_and_actions(
  )
  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
+      tool_name, tool_instructions
  )
  mock_openapi_action_spec_parser.return_value.parse.assert_called_once()
  assert len(toolset.get_tools()) == 1
@@ -390,6 +389,5 @@ def test_initialization_with_connection_details(
      tool_instructions=tool_instructions,
  )
  mock_integration_client.return_value.get_openapi_spec_for_connection.assert_called_once_with(
-      tool_name,
-      tool_instructions
+      tool_name, tool_instructions
  )
--- a/tests/unittests/tools/retrieval/init.py
+++ b/tests/unittests/tools/retrieval/init.py
@@ -11,4 +11,3 @@
 # 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.
-