dify

dify/api/tests/unit_tests/tools/test_api_tool.py

@@ -0,0 +1,249 @@
+import json
+import operator
+from typing import TypeVar
+from unittest.mock import Mock, patch
+
+import httpx
+import pytest
+
+from core.tools.__base.tool_runtime import ToolRuntime
+from core.tools.custom_tool.tool import ApiTool
+from core.tools.entities.common_entities import I18nObject
+from core.tools.entities.tool_bundle import ApiToolBundle
+from core.tools.entities.tool_entities import (
+    ToolEntity,
+    ToolIdentity,
+    ToolInvokeMessage,
+)
+
+_T = TypeVar("_T")
+
+
+def _get_message_by_type(msgs: list[ToolInvokeMessage], msg_type: type[_T]) -> ToolInvokeMessage | None:
+    return next((i for i in msgs if isinstance(i.message, msg_type)), None)
+
+
+class TestApiToolInvoke:
+    """Test suite for ApiTool._invoke method to ensure JSON responses are properly serialized."""
+
+    def setup_method(self):
+        """Setup test fixtures."""
+        # Create a mock tool entity
+        self.mock_tool_identity = ToolIdentity(
+            author="test",
+            name="test_api_tool",
+            label=I18nObject(en_US="Test API Tool", zh_Hans="测试API工具"),
+            provider="test_provider",
+        )
+        self.mock_tool_entity = ToolEntity(identity=self.mock_tool_identity)
+
+        # Create a mock API bundle
+        self.mock_api_bundle = ApiToolBundle(
+            server_url="https://api.example.com/test",
+            method="GET",
+            openapi={},
+            operation_id="test_operation",
+            parameters=[],
+            author="test_author",
+        )
+
+        # Create a mock runtime
+        self.mock_runtime = Mock(spec=ToolRuntime)
+        self.mock_runtime.credentials = {"auth_type": "none"}
+
+        # Create the ApiTool instance
+        self.api_tool = ApiTool(
+            entity=self.mock_tool_entity,
+            api_bundle=self.mock_api_bundle,
+            runtime=self.mock_runtime,
+            provider_id="test_provider",
+        )
+
+    @patch("core.tools.custom_tool.tool.ssrf_proxy.get")
+    def test_invoke_with_json_response_creates_text_message_with_serialized_json(self, mock_get: Mock) -> None:
+        """Test that when upstream returns JSON, the output Text message contains JSON-serialized string."""
+        # Setup mock response with JSON content
+        json_response_data = {
+            "key": "value",
+            "number": 123,
+            "nested": {"inner": "data"},
+        }
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.content = json.dumps(json_response_data).encode("utf-8")
+        mock_response.json.return_value = json_response_data
+        mock_response.text = json.dumps(json_response_data)
+        mock_response.headers = {"content-type": "application/json"}
+        mock_get.return_value = mock_response
+
+        # Invoke the tool
+        result_generator = self.api_tool._invoke(user_id="test_user", tool_parameters={})
+
+        # Get the result from the generator
+        result = list(result_generator)
+        assert len(result) == 2
+
+        # Verify _invoke yields text message
+        text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage)
+        assert text_message is not None, "_invoke should yield a text message"
+        assert isinstance(text_message, ToolInvokeMessage)
+        assert text_message.type == ToolInvokeMessage.MessageType.TEXT
+        assert text_message.message is not None
+        # Verify the text contains the JSON-serialized string
+        # Check if message is a TextMessage
+        assert isinstance(text_message.message, ToolInvokeMessage.TextMessage)
+        # Verify it's a valid JSON string and equals to the mock response
+        parsed_back = json.loads(text_message.message.text)
+        assert parsed_back == json_response_data
+
+        # Verify _invoke yields json message
+        json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage)
+        assert json_message is not None, "_invoke should yield a JSON message"
+        assert isinstance(json_message, ToolInvokeMessage)
+        assert json_message.type == ToolInvokeMessage.MessageType.JSON
+        assert json_message.message is not None
+
+        assert isinstance(json_message.message, ToolInvokeMessage.JsonMessage)
+        assert json_message.message.json_object == json_response_data
+
+    @patch("core.tools.custom_tool.tool.ssrf_proxy.get")
+    @pytest.mark.parametrize(
+        "test_case",
+        [
+            (
+                "array",
+                [
+                    {"id": 1, "name": "Item 1", "active": True},
+                    {"id": 2, "name": "Item 2", "active": False},
+                    {"id": 3, "name": "项目 3", "active": True},
+                ],
+            ),
+            (
+                "string",
+                "string",
+            ),
+            (
+                "number",
+                123.456,
+            ),
+            (
+                "boolean",
+                True,
+            ),
+            (
+                "null",
+                None,
+            ),
+        ],
+        ids=operator.itemgetter(0),
+    )
+    def test_invoke_with_non_dict_json_response_creates_text_message_with_serialized_json(
+        self, mock_get: Mock, test_case
+    ) -> None:
+        """Test that when upstream returns a non-dict JSON, the output Text message contains JSON-serialized string."""
+        # Setup mock response with non-dict JSON content
+        _, json_value = test_case
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.content = json.dumps(json_value).encode("utf-8")
+        mock_response.json.return_value = json_value
+        mock_response.text = json.dumps(json_value)
+        mock_response.headers = {"content-type": "application/json"}
+        mock_get.return_value = mock_response
+
+        # Invoke the tool
+        result_generator = self.api_tool._invoke(user_id="test_user", tool_parameters={})
+
+        # Get the result from the generator
+        result = list(result_generator)
+        assert len(result) == 1
+
+        # Verify  _invoke yields a text message
+        text_message = _get_message_by_type(result, ToolInvokeMessage.TextMessage)
+        assert text_message is not None, "_invoke should yield a text message containing the serialized JSON."
+        assert isinstance(text_message, ToolInvokeMessage)
+        assert text_message.type == ToolInvokeMessage.MessageType.TEXT
+        assert text_message.message is not None
+        # Verify the text contains the JSON-serialized string
+        # Check if message is a TextMessage
+        assert isinstance(text_message.message, ToolInvokeMessage.TextMessage)
+        # Verify it's a valid JSON string
+        parsed_back = json.loads(text_message.message.text)
+        assert parsed_back == json_value
+
+        # Verify _invoke yields json message
+        json_message = _get_message_by_type(result, ToolInvokeMessage.JsonMessage)
+        assert json_message is None, "_invoke should not yield a JSON message for JSON array response"
+
+    @patch("core.tools.custom_tool.tool.ssrf_proxy.get")
+    def test_invoke_with_text_response_creates_text_message_with_original_text(self, mock_get: Mock) -> None:
+        """Test that when upstream returns plain text, the output Text message contains the original text."""
+        # Setup mock response with plain text content
+        text_response_data = "This is a plain text response"
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.content = text_response_data.encode("utf-8")
+        mock_response.json.side_effect = json.JSONDecodeError("Expecting value", "doc", 0)
+        mock_response.text = text_response_data
+        mock_response.headers = {"content-type": "text/plain"}
+        mock_get.return_value = mock_response
+
+        # Invoke the tool
+        result_generator = self.api_tool._invoke(user_id="test_user", tool_parameters={})
+
+        # Get the result from the generator
+        result = list(result_generator)
+        assert len(result) == 1
+
+        # Verify it's a text message with the original text
+        message = result[0]
+        assert isinstance(message, ToolInvokeMessage)
+        assert message.type == ToolInvokeMessage.MessageType.TEXT
+        assert message.message is not None
+        # Check if message is a TextMessage
+        assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+        assert message.message.text == text_response_data
+
+    @patch("core.tools.custom_tool.tool.ssrf_proxy.get")
+    def test_invoke_with_empty_response(self, mock_get: Mock) -> None:
+        """Test that empty responses are handled correctly."""
+        # Setup mock response with empty content
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.content = b""
+        mock_response.headers = {"content-type": "application/json"}
+        mock_get.return_value = mock_response
+
+        # Invoke the tool
+        result_generator = self.api_tool._invoke(user_id="test_user", tool_parameters={})
+
+        # Get the result from the generator
+        result = list(result_generator)
+        assert len(result) == 1
+
+        # Verify it's a text message with the empty response message
+        message = result[0]
+        assert isinstance(message, ToolInvokeMessage)
+        assert message.type == ToolInvokeMessage.MessageType.TEXT
+        assert message.message is not None
+        # Check if message is a TextMessage
+        assert isinstance(message.message, ToolInvokeMessage.TextMessage)
+        assert "Empty response from the tool" in message.message.text
+
+    @patch("core.tools.custom_tool.tool.ssrf_proxy.get")
+    def test_invoke_with_error_response(self, mock_get: Mock) -> None:
+        """Test that error responses are handled correctly."""
+        # Setup mock response with error status code
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.status_code = 404
+        mock_response.text = "Not Found"
+        mock_get.return_value = mock_response
+
+        result_generator = self.api_tool._invoke(user_id="test_user", tool_parameters={})
+
+        # Invoke the tool and expect an error
+        with pytest.raises(Exception) as exc_info:
+            list(result_generator)  # Consume the generator to trigger the error
+
+        # Verify the error message
+        assert "Request failed with status code 404" in str(exc_info.value)