98 lines
4.5 KiB
Python
98 lines
4.5 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
from api import chat
|
|
from graph.session_manager import SessionManager
|
|
from graph.chat_model_options import chat_model_options
|
|
|
|
|
|
class ChatRepairTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.sessions = SessionManager()
|
|
self.sessions.initialize(Path(self.temp.name))
|
|
self.patch = patch.object(chat,"session_manager",self.sessions)
|
|
self.patch.start()
|
|
|
|
async def asyncTearDown(self):
|
|
self.patch.stop()
|
|
self.temp.cleanup()
|
|
|
|
async def run_events(self, events, error=None):
|
|
async def fake(*args):
|
|
for item in events:
|
|
yield item
|
|
if error:
|
|
raise error
|
|
with patch.object(chat.agent_manager,"astream",fake), patch.object(chat,"_generate_title",AsyncMock(return_value=None)):
|
|
return [{"event":e["event"],"data":json.loads(e["data"])} async for e in chat.event_generator("测试", "chat-test")]
|
|
|
|
async def test_empty_answer_reports_error_and_never_saves_blank(self):
|
|
result=await self.run_events([{"type":"done","content":""}])
|
|
self.assertEqual(result[-1]["data"]["code"],"empty_response")
|
|
self.assertNotIn("done",[e["event"] for e in result])
|
|
messages=self.sessions.load_session("chat-test")
|
|
self.assertTrue(messages[-1]["content"].strip())
|
|
self.assertIn("未完成",messages[-1]["content"])
|
|
|
|
async def test_final_without_tokens_becomes_visible(self):
|
|
result=await self.run_events([{"type":"done","content":"有效回答"}])
|
|
self.assertEqual([e["event"] for e in result],["token","done"])
|
|
self.assertEqual(self.sessions.load_session("chat-test")[-1]["content"],"有效回答")
|
|
|
|
async def test_tokens_are_not_duplicated_and_tool_segments_preserved(self):
|
|
result=await self.run_events([{"type":"token","content":"先查"},{"type":"tool_start","tool":"read","input":"x"},{"type":"tool_end","tool":"read","output":"y"},{"type":"new_response"},{"type":"token","content":"结果"},{"type":"done","content":"先查结果"}])
|
|
self.assertEqual([e["data"]["content"] for e in result if e["event"]=="token"],["先查","结果"])
|
|
messages=self.sessions.load_session("chat-test")
|
|
self.assertEqual(len(messages),3)
|
|
self.assertEqual(messages[-1]["content"],"结果")
|
|
|
|
async def test_failure_keeps_partial_and_hides_credentials(self):
|
|
result=await self.run_events([{"type":"token","content":"部分结果"}],RuntimeError("secret-key=private"))
|
|
self.assertEqual(result[-1]["data"]["code"],"chat_failed")
|
|
self.assertNotIn("secret-key",json.dumps(result))
|
|
self.assertIn("部分结果",self.sessions.load_session("chat-test")[-1]["content"])
|
|
|
|
async def test_timeout_and_incomplete_stream(self):
|
|
result=await self.run_events([],TimeoutError())
|
|
self.assertEqual(result[-1]["data"]["code"],"timeout")
|
|
result=await self.run_events([])
|
|
self.assertEqual(result[-1]["data"]["code"],"chat_failed")
|
|
|
|
async def test_bounded_generator_is_cancelled_without_replay(self):
|
|
calls=[]
|
|
async def stalled():
|
|
try:
|
|
calls.append("start")
|
|
await asyncio.sleep(1)
|
|
yield {"type":"done"}
|
|
finally:
|
|
calls.append("closed")
|
|
with self.assertRaises(TimeoutError):
|
|
async for _ in chat.bounded_events(stalled(),idle_seconds=.01,total_seconds=.02):
|
|
pass
|
|
self.assertEqual(calls,["start","closed"])
|
|
|
|
def test_chat_options_do_not_mutate_scoped_models(self):
|
|
with patch.dict(os.environ,{"CHAT_THINKING_MODE":"disabled"}):
|
|
self.assertEqual(chat_model_options()["extra_body"],{"thinking":{"type":"disabled"}})
|
|
with patch.dict(os.environ,{"CHAT_THINKING_MODE":""}):
|
|
self.assertNotIn("extra_body",chat_model_options())
|
|
with patch.dict(os.environ,{"CHAT_THINKING_MODE":"invalid"}):
|
|
with self.assertRaises(ValueError):chat_model_options()
|
|
|
|
def test_frontend_streaming_and_proxy(self):
|
|
root=Path(__file__).resolve().parents[1]
|
|
result=subprocess.run([shutil.which("node") or "node","scripts/test_chat_stream.cjs"],cwd=root,capture_output=True,text=True,encoding="utf-8",timeout=25)
|
|
self.assertEqual(result.returncode,0,result.stdout+result.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|