fix application integration toolset unit tests

PiperOrigin-RevId: 758080966

Author: seanzhougoogle

tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py

@@ -47,7 +47,14 @@ def mock_openapi_toolset():
     mock_toolset_instance = mock.MagicMock()
     mock_rest_api_tool = mock.MagicMock(spec=rest_api_tool.RestApiTool)
     mock_rest_api_tool.name = "Test Tool"
-    mock_toolset_instance.get_tools.return_value = [mock_rest_api_tool]
+
+    # Create an async mock for the get_tools method
+    async def mock_get_tools():
+      return [mock_rest_api_tool]
+
+    # Assign the async mock function to get_tools
+    mock_toolset_instance.get_tools = mock_get_tools
+
     mock_toolset.return_value = mock_toolset_instance
     yield mock_toolset
 
@@ -62,10 +69,12 @@ def mock_openapi_toolset_with_multiple_tools_and_no_tools():
     mock_rest_api_tool.name = "Test Tool"
     mock_rest_api_tool_2 = mock.MagicMock(spec=rest_api_tool.RestApiTool)
     mock_rest_api_tool_2.name = "Test Tool 2"
-    mock_toolset_instance.get_tools.return_value = [
-        mock_rest_api_tool,
-        mock_rest_api_tool_2,
-    ]
+
+    # Create an async mock for the get_tools method
+    async def mock_get_tools():
+      return [mock_rest_api_tool, mock_rest_api_tool_2]
+
+    mock_toolset_instance.get_tools = mock_get_tools
     mock_toolset.return_value = mock_toolset_instance
     yield mock_toolset
 
@@ -152,7 +161,8 @@ def connection_details():
   }
 
 
-def test_initialization_with_integration_and_trigger(
+@pytest.mark.asyncio
+async def test_initialization_with_integration_and_trigger(
     project,
     location,
     mock_integration_client,
@@ -170,11 +180,13 @@ def test_initialization_with_integration_and_trigger(
   mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
   mock_connections_client.assert_not_called()
   mock_openapi_toolset.assert_called_once()
-  assert len(toolset.get_tools()) == 1
-  assert toolset.get_tools()[0].name == "Test Tool"
+  tools = await toolset.get_tools()
+  assert len(tools) == 1
+  assert tools[0].name == "Test Tool"
 
 
-def test_initialization_with_integration_and_list_of_triggers(
+@pytest.mark.asyncio
+async def test_initialization_with_integration_and_list_of_triggers(
     project,
     location,
     mock_integration_client,
@@ -199,12 +211,14 @@ def test_initialization_with_integration_and_list_of_triggers(
   mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
   mock_connections_client.assert_not_called()
   mock_openapi_toolset_with_multiple_tools_and_no_tools.assert_called_once()
-  assert len(toolset.get_tools()) == 2
-  assert toolset.get_tools()[0].name == "Test Tool"
-  assert toolset.get_tools()[1].name == "Test Tool 2"
+  tools = await toolset.get_tools()
+  assert len(tools) == 2
+  assert tools[0].name == "Test Tool"
+  assert tools[1].name == "Test Tool 2"
 
 
-def test_initialization_with_integration_and_empty_trigger_list(
+@pytest.mark.asyncio
+async def test_initialization_with_integration_and_empty_trigger_list(
     project,
     location,
     mock_integration_client,
@@ -221,12 +235,14 @@ def test_initialization_with_integration_and_empty_trigger_list(
   mock_integration_client.return_value.get_openapi_spec_for_integration.assert_called_once()
   mock_connections_client.assert_not_called()
   mock_openapi_toolset_with_multiple_tools_and_no_tools.assert_called_once()
-  assert len(toolset.get_tools()) == 2
-  assert toolset.get_tools()[0].name == "Test Tool"
-  assert toolset.get_tools()[1].name == "Test Tool 2"
+  tools = await toolset.get_tools()
+  assert len(tools) == 2
+  assert tools[0].name == "Test Tool"
+  assert tools[1].name == "Test Tool 2"
 
 
-def test_initialization_with_connection_and_entity_operations(
+@pytest.mark.asyncio
+async def test_initialization_with_connection_and_entity_operations(
     project,
     location,
     mock_integration_client,
@@ -268,14 +284,17 @@ def test_initialization_with_connection_and_entity_operations(
       tool_name,
       tool_instructions,
   )
-  assert len(toolset.get_tools()) == 1
-  assert toolset.get_tools()[0].name == "list_issues"
-  assert isinstance(toolset.get_tools()[0], IntegrationConnectorTool)
-  assert toolset.get_tools()[0].entity == "Issues"
-  assert toolset.get_tools()[0].operation == "LIST_ENTITIES"
+
+  tools = await toolset.get_tools()
+  assert len(tools) == 1
+  assert tools[0].name == "list_issues"
+  assert isinstance(tools[0], IntegrationConnectorTool)
+  assert tools[0].entity == "Issues"
+  assert tools[0].operation == "LIST_ENTITIES"
 
 
-def test_initialization_with_connection_and_actions(
+@pytest.mark.asyncio
+async def test_initialization_with_connection_and_actions(
     project,
     location,
     mock_integration_client,
@@ -309,11 +328,12 @@ def test_initialization_with_connection_and_actions(
       tool_name, tool_instructions
   )
   mock_openapi_action_spec_parser.return_value.parse.assert_called_once()
-  assert len(toolset.get_tools()) == 1
-  assert toolset.get_tools()[0].name == "list_issues_operation"
-  assert isinstance(toolset.get_tools()[0], IntegrationConnectorTool)
-  assert toolset.get_tools()[0].action == "CustomAction"
-  assert toolset.get_tools()[0].operation == "EXECUTE_ACTION"
+  tools = await toolset.get_tools()
+  assert len(tools) == 1
+  assert tools[0].name == "list_issues_operation"
+  assert isinstance(tools[0], IntegrationConnectorTool)
+  assert tools[0].action == "CustomAction"
+  assert tools[0].operation == "EXECUTE_ACTION"
 
 
 def test_initialization_without_required_params(project, location):
@@ -412,15 +432,16 @@ def test_initialization_without_explicit_service_account_credentials(
   assert kwargs["auth_credential"].service_account.use_default_credential
 
 
-def test_get_tools(
+@pytest.mark.asyncio
+async def test_get_tools(
     project, location, mock_integration_client, mock_openapi_toolset
 ):
   integration_name = "test-integration"
   triggers = ["test-trigger"]
   toolset = ApplicationIntegrationToolset(
       project, location, integration=integration_name, triggers=triggers
   )
-  tools = toolset.get_tools()
+  tools = await toolset.get_tools()
   assert len(tools) == 1
   assert isinstance(tools[0], rest_api_tool.RestApiTool)
   assert tools[0].name == "Test Tool"