mirror of
https://github.com/langgenius/dify.git
synced 2026-07-22 03:08:33 +08:00
261 lines
11 KiB
Python
261 lines
11 KiB
Python
import json
|
|
import operator
|
|
from collections.abc import Iterator
|
|
from unittest.mock import Mock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from sqlalchemy import Engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def database_session(sqlite_engine: Engine) -> Iterator[Session]:
|
|
"""Yield the live ORM session required by the shared tool invocation contract."""
|
|
with Session(sqlite_engine) as session:
|
|
yield session
|
|
|
|
|
|
def _get_message_by_type[T](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, database_session: Session
|
|
) -> 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(session=database_session, 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, database_session: Session
|
|
) -> 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(session=database_session, 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, database_session: Session
|
|
) -> 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(session=database_session, 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, database_session: Session) -> 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(session=database_session, 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, database_session: Session) -> 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(session=database_session, 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)
|