|
| 1 | +import unittest |
| 2 | +from unittest.mock import MagicMock, patch |
| 3 | + |
| 4 | +from fastdeploy.input.ernie_processor import ErnieProcessor |
| 5 | + |
| 6 | + |
| 7 | +class TestErnieProcessorProcessResponseDictStreaming(unittest.TestCase): |
| 8 | + def setUp(self): |
| 9 | + # 创建 ErnieProcessor 实例的模拟对象 |
| 10 | + with patch.object(ErnieProcessor, "__init__", return_value=None) as mock_init: |
| 11 | + self.processor = ErnieProcessor("model_path") |
| 12 | + mock_init.side_effect = lambda *args, **kwargs: print(f"__init__ called with {args}, {kwargs}") |
| 13 | + |
| 14 | + # 设置必要的属性 |
| 15 | + self.processor.tokenizer = MagicMock() |
| 16 | + self.processor.tokenizer.eos_token_id = 1 |
| 17 | + self.processor.decode_status = {} |
| 18 | + self.processor.tool_parsers = {} |
| 19 | + |
| 20 | + # 模拟 ids2tokens 方法 |
| 21 | + def mock_ids2tokens(token_ids, task_id): |
| 22 | + return "delta_text", [2, 3], "previous_texts" |
| 23 | + |
| 24 | + self.processor.ids2tokens = mock_ids2tokens |
| 25 | + |
| 26 | + # 模拟推理解析器 |
| 27 | + self.mock_reasoning_parser = MagicMock() |
| 28 | + self.mock_reasoning_parser.__class__.__name__ = "ErnieX1ReasoningParser" |
| 29 | + self.mock_reasoning_parser.extract_reasoning_content_streaming.return_value = ("reasoning", "text") |
| 30 | + self.processor.reasoning_parser = self.mock_reasoning_parser |
| 31 | + |
| 32 | + # 模拟工具解析器 |
| 33 | + self.mock_tool_parser = MagicMock() |
| 34 | + self.mock_tool_parser.extract_tool_calls_streaming.return_value = "tool_call" |
| 35 | + self.mock_tool_parser_obj = MagicMock() |
| 36 | + self.mock_tool_parser_obj.return_value = self.mock_tool_parser |
| 37 | + self.processor.tool_parser_obj = self.mock_tool_parser_obj |
| 38 | + |
| 39 | + def test_process_response_dict_streaming_normal_case(self): |
| 40 | + """测试正常情况下的流式响应处理""" |
| 41 | + # 准备输入 |
| 42 | + response_dict = {"finished": False, "request_id": "req1", "outputs": {"token_ids": [4, 5]}} |
| 43 | + kwargs = {"enable_thinking": True} |
| 44 | + |
| 45 | + # 调用方法 |
| 46 | + result = self.processor.process_response_dict_streaming(response_dict, **kwargs) |
| 47 | + |
| 48 | + # 验证结果 |
| 49 | + self.assertEqual(result["outputs"]["raw_prediction"], "delta_text") |
| 50 | + |
| 51 | + |
| 52 | +if __name__ == "__main__": |
| 53 | + unittest.main() |
0 commit comments