cnbugs
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
import asyncio
|
|
import json
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import backend.main as main
|
|
|
|
|
|
def parse_sse(event):
|
|
assert event.startswith("data: ")
|
|
return json.loads(event[6:].strip())
|
|
|
|
|
|
class AnalyzeProcessStreamTest(unittest.TestCase):
|
|
def test_analysis_stream_emits_visible_process_steps_before_content(self):
|
|
async def fake_stream(report, question=""):
|
|
yield "分析结论"
|
|
|
|
async def collect_events():
|
|
main._last_diagnosis = {"score": 80}
|
|
main._last_report = "diagnosis report"
|
|
with patch.object(main, "analyze_with_ai_stream", fake_stream):
|
|
response = await main.analyze_stream(question="关注严重问题")
|
|
return [parse_sse(event) async for event in response.body_iterator]
|
|
|
|
events = asyncio.run(collect_events())
|
|
process_events = [event for event in events if event["type"] == "process"]
|
|
|
|
self.assertEqual(
|
|
[event["status"] for event in process_events],
|
|
["running", "done", "running", "done", "running", "done", "running", "done"],
|
|
)
|
|
self.assertEqual(
|
|
[event["title"] for event in process_events[::2]],
|
|
["读取诊断报告", "识别关键风险", "生成修复方案", "整理分析结论"],
|
|
)
|
|
first_chunk_index = next(i for i, event in enumerate(events) if event["type"] == "chunk")
|
|
final_process_index = max(i for i, event in enumerate(events) if event["type"] == "process")
|
|
self.assertLess(final_process_index, first_chunk_index)
|
|
self.assertEqual(events[-1]["type"], "done")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|