diff --git a/src/google/adk/tools/application_integration_tool/clients/connections_client.py b/src/google/adk/tools/application_integration_tool/clients/connections_client.py index 0a1b6dbc42d..749c56a733c 100644 --- a/src/google/adk/tools/application_integration_tool/clients/connections_client.py +++ b/src/google/adk/tools/application_integration_tool/clients/connections_client.py @@ -41,6 +41,8 @@ "integrations.mtls.googleapis.com" ) _DEFAULT_REQUEST_TIMEOUT_SECONDS = 30 +# Upper bound on waiting for a schema long-running operation to finish. +_DEFAULT_OPERATION_TIMEOUT_SECONDS = 300 class _ServiceAccountCredentialsFactory(Protocol): @@ -962,15 +964,29 @@ def _poll_operation(self, operation_id: str) -> Dict[str, Any]: Raises: PermissionError: If there are credential issues. - ValueError: If there's a request error. + ValueError: If there's a request error, or the operation finished with + an error. + TimeoutError: If the operation is not done within + `_DEFAULT_OPERATION_TIMEOUT_SECONDS`. Exception: For any other unexpected errors. """ - operation_done: bool = False - operation_response: Dict[str, Any] = {} - while not operation_done: - get_operation_url = f"{self.connector_url}/v1/{operation_id}" + get_operation_url = f"{self.connector_url}/v1/{operation_id}" + deadline = time.monotonic() + _DEFAULT_OPERATION_TIMEOUT_SECONDS + while True: response = self._execute_api_call(get_operation_url) operation_response = self._response_json(response) - operation_done = bool(operation_response.get("done", False)) + if operation_response.get("done", False): + break + if time.monotonic() >= deadline: + raise TimeoutError( + f"Operation {operation_id} did not finish within" + f" {_DEFAULT_OPERATION_TIMEOUT_SECONDS} seconds." + ) time.sleep(1) + # A finished operation carries either `response` or `error`; without this + # check a failed lookup reads as an empty schema and yields no tools. + error = operation_response.get("error") + if error: + message = error.get("message") if isinstance(error, dict) else None + raise ValueError(f"Operation {operation_id} failed: {message or error}") return operation_response diff --git a/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py b/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py index 770847a6f18..ee317d3c238 100644 --- a/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py +++ b/tests/unittests/tools/application_integration_tool/clients/test_connections_client.py @@ -360,6 +360,90 @@ def test_get_entity_schema_and_operations_execute_api_call_error( with pytest.raises(ValueError, match="Request error"): client.get_entity_schema_and_operations("entity1") + def test_get_entity_schema_and_operations_failed_operation( + self, project, location, connection_name + ): + """A failed operation must not read as an empty schema with no tools.""" + credentials = {"email": "test@example.com"} + client = ConnectionsClient(project, location, connection_name, credentials) + mock_execute_response_initial = mock.MagicMock() + mock_execute_response_initial.json.return_value = { + "name": "operations/test_op" + } + mock_execute_response_poll_failed = mock.MagicMock() + mock_execute_response_poll_failed.json.return_value = { + "done": True, + "error": {"code": 5, "message": "Entity type 'entity1' not found."}, + } + + with mock.patch.object( + client, + "_execute_api_call", + side_effect=[ + mock_execute_response_initial, + mock_execute_response_poll_failed, + ], + ): + with pytest.raises( + ValueError, + match=( + "Operation operations/test_op failed: Entity type 'entity1' not" + " found." + ), + ): + client.get_entity_schema_and_operations("entity1") + + def test_poll_operation_times_out(self, project, location, connection_name): + credentials = {"email": "test@example.com"} + client = ConnectionsClient(project, location, connection_name, credentials) + mock_execute_response_pending = mock.MagicMock() + mock_execute_response_pending.json.return_value = {"done": False} + + with ( + mock.patch.object( + client, + "_execute_api_call", + return_value=mock_execute_response_pending, + ), + mock.patch( + "google.adk.tools.application_integration_tool.clients.connections_client.time" + ) as mock_time, + ): + mock_time.monotonic.side_effect = [0, 1, 301] + with pytest.raises( + TimeoutError, + match="Operation operations/test_op did not finish within 300", + ): + client._poll_operation("operations/test_op") + assert client._execute_api_call.call_count == 2 + mock_time.sleep.assert_called_once_with(1) + + def test_poll_operation_does_not_sleep_once_done( + self, project, location, connection_name + ): + credentials = {"email": "test@example.com"} + client = ConnectionsClient(project, location, connection_name, credentials) + mock_execute_response_done = mock.MagicMock() + mock_execute_response_done.json.return_value = { + "done": True, + "response": {"jsonSchema": {}}, + } + + with ( + mock.patch.object( + client, "_execute_api_call", return_value=mock_execute_response_done + ), + mock.patch( + "google.adk.tools.application_integration_tool.clients.connections_client.time" + ) as mock_time, + ): + mock_time.monotonic.return_value = 0 + assert client._poll_operation("operations/test_op") == { + "done": True, + "response": {"jsonSchema": {}}, + } + mock_time.sleep.assert_not_called() + def test_get_action_schema_success( self, project, location, connection_name, mock_credentials ):