"""
|
测试Dify流式模式功能
|
"""
|
|
import pytest
|
import json
|
from unittest.mock import Mock, patch, MagicMock
|
from app.services.dify_client import DifyClient
|
from config import settings
|
|
|
class TestDifyStreaming:
|
"""测试Dify流式模式"""
|
|
def setup_method(self):
|
"""测试前设置"""
|
self.client = DifyClient()
|
|
def test_process_stream_response_success(self):
|
"""测试成功处理流式响应"""
|
# 模拟流式响应数据
|
stream_data = [
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \"Hello\", \"created_at\": 1705398420}",
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \" World\", \"created_at\": 1705398420}",
|
"data: {\"event\": \"message_end\", \"metadata\": {\"usage\": {\"total_tokens\": 10}}, \"usage\": {\"total_tokens\": 10}}",
|
]
|
|
# 创建模拟响应对象
|
mock_response = Mock()
|
mock_response.iter_lines.return_value = stream_data
|
|
# 测试处理流式响应
|
result = self.client._process_stream_response(mock_response, "test_user")
|
|
# 验证结果
|
assert result is not None
|
assert result["answer"] == "Hello World"
|
assert result["conversation_id"] == "test-conv"
|
assert result["task_id"] == "test-task"
|
assert result["usage"]["total_tokens"] == 10
|
|
def test_process_stream_response_error(self):
|
"""测试处理流式响应错误"""
|
# 模拟错误响应数据
|
stream_data = [
|
"data: {\"event\": \"error\", \"message\": \"API调用失败\", \"code\": \"500\"}",
|
]
|
|
# 创建模拟响应对象
|
mock_response = Mock()
|
mock_response.iter_lines.return_value = stream_data
|
|
# 测试处理流式响应
|
result = self.client._process_stream_response(mock_response, "test_user")
|
|
# 验证结果
|
assert result is None
|
|
def test_process_stream_response_incomplete(self):
|
"""测试处理不完整的流式响应"""
|
# 模拟不完整响应数据（缺少message_end事件）
|
stream_data = [
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \"Hello\", \"created_at\": 1705398420}",
|
]
|
|
# 创建模拟响应对象
|
mock_response = Mock()
|
mock_response.iter_lines.return_value = stream_data
|
|
# 测试处理流式响应
|
result = self.client._process_stream_response(mock_response, "test_user")
|
|
# 验证结果 - 即使没有message_end事件，只要有内容和conversation_id也应该返回结果
|
assert result is not None
|
assert result["answer"] == "Hello"
|
assert result["conversation_id"] == "test-conv"
|
|
@patch('app.services.dify_client.settings')
|
def test_send_message_uses_streaming_when_enabled(self, mock_settings):
|
"""测试当启用流式模式时使用流式发送"""
|
# 设置配置为启用流式模式
|
mock_settings.dify_streaming_enabled = True
|
|
# 模拟流式发送方法
|
with patch.object(self.client, 'send_chat_message_stream') as mock_stream:
|
mock_stream.return_value = {"answer": "test response", "conversation_id": "test-conv"}
|
|
result = self.client.send_message("test query", "test_user")
|
|
# 验证调用了流式方法
|
mock_stream.assert_called_once_with("test query", "test_user", None, None)
|
assert result["answer"] == "test response"
|
|
@patch('app.services.dify_client.settings')
|
def test_send_message_uses_blocking_when_disabled(self, mock_settings):
|
"""测试当禁用流式模式时使用阻塞发送"""
|
# 设置配置为禁用流式模式
|
mock_settings.dify_streaming_enabled = False
|
|
# 模拟阻塞发送方法
|
with patch.object(self.client, 'send_chat_message') as mock_blocking:
|
mock_blocking.return_value = {"answer": "test response", "conversation_id": "test-conv"}
|
|
result = self.client.send_message("test query", "test_user")
|
|
# 验证调用了阻塞方法
|
mock_blocking.assert_called_once_with("test query", "test_user", None, None)
|
assert result["answer"] == "test response"
|
|
@patch('app.services.dify_client.settings')
|
def test_send_message_force_streaming_override(self, mock_settings):
|
"""测试强制流式模式覆盖配置"""
|
# 设置配置为禁用流式模式
|
mock_settings.dify_streaming_enabled = False
|
|
# 模拟流式发送方法
|
with patch.object(self.client, 'send_chat_message_stream') as mock_stream:
|
mock_stream.return_value = {"answer": "test response", "conversation_id": "test-conv"}
|
|
# 强制使用流式模式
|
result = self.client.send_message("test query", "test_user", force_streaming=True)
|
|
# 验证调用了流式方法（覆盖了配置）
|
mock_stream.assert_called_once_with("test query", "test_user", None, None)
|
assert result["answer"] == "test response"
|
|
def test_process_stream_response_with_ping_events(self):
|
"""测试处理包含ping事件的流式响应"""
|
# 模拟包含ping事件的响应数据
|
stream_data = [
|
"data: {\"event\": \"ping\"}",
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \"Hello\", \"created_at\": 1705398420}",
|
"data: {\"event\": \"ping\"}",
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \" World\", \"created_at\": 1705398420}",
|
"data: {\"event\": \"message_end\", \"metadata\": {\"usage\": {\"total_tokens\": 10}}}",
|
]
|
|
# 创建模拟响应对象
|
mock_response = Mock()
|
mock_response.iter_lines.return_value = stream_data
|
|
# 测试处理流式响应
|
result = self.client._process_stream_response(mock_response, "test_user")
|
|
# 验证结果（ping事件应该被忽略）
|
assert result is not None
|
assert result["answer"] == "Hello World"
|
assert result["conversation_id"] == "test-conv"
|
|
def test_process_stream_response_with_invalid_json(self):
|
"""测试处理包含无效JSON的流式响应"""
|
# 模拟包含无效JSON的响应数据
|
stream_data = [
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \"Hello\", \"created_at\": 1705398420}",
|
"data: invalid json data",
|
"data: {\"event\": \"message\", \"task_id\": \"test-task\", \"id\": \"test-msg\", \"conversation_id\": \"test-conv\", \"answer\": \" World\", \"created_at\": 1705398420}",
|
"data: {\"event\": \"message_end\"}",
|
]
|
|
# 创建模拟响应对象
|
mock_response = Mock()
|
mock_response.iter_lines.return_value = stream_data
|
|
# 测试处理流式响应
|
result = self.client._process_stream_response(mock_response, "test_user")
|
|
# 验证结果（无效JSON应该被跳过）
|
assert result is not None
|
assert result["answer"] == "Hello World"
|
assert result["conversation_id"] == "test-conv"
|
|
|
if __name__ == "__main__":
|
pytest.main([__file__])
|
