From dd8e9b71e525ae149742d4f6330d84f293fe8562 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Tue, 13 May 2025 00:06:23 -0700 Subject: [PATCH] fix application integration toolset unit tests PiperOrigin-RevId: 758080966 --- .../test_application_integration_toolset.py | 81 ++++++++++++------- 1 file changed, 51 insertions(+), 30 deletions(-) diff --git a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py index 0a707e6..b2f57c2 100644 --- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py +++ b/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,7 +432,8 @@ 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" @@ -420,7 +441,7 @@ def test_get_tools( 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"