from __future__ import annotations import io import tempfile import unittest from pathlib import Path from unittest.mock import patch from assistant_agent import ReviewAssistantError, stream_review_assistant from database import ReviewDatabase from server import RequestHandler class StreamingResponse: def __init__(self, lines: list[bytes]): self.lines = lines def __enter__(self): return iter(self.lines) def __exit__(self, exc_type, exc, traceback): return False class ReviewAssistantStreamingTests(unittest.TestCase): def test_openai_compatible_sse_is_yielded_in_order(self): response = StreamingResponse( [ b'data: {"choices":[{"delta":{"content":"first"}}]}\n', b'data: {"choices":[{"delta":{"content":" second"}}]}\n', b'data: [DONE]\n', ] ) with patch("assistant_agent.urllib.request.urlopen", return_value=response): chunks = list( stream_review_assistant({}, "question", [], "key", "https://example.test/v1", "model") ) self.assertEqual(chunks, ["first", " second"]) def test_empty_stream_is_rejected(self): response = StreamingResponse([b"data: [DONE]\n"]) with patch("assistant_agent.urllib.request.urlopen", return_value=response): with self.assertRaises(ReviewAssistantError): list( stream_review_assistant({}, "question", [], "key", "https://example.test/v1", "model") ) def test_final_full_message_does_not_duplicate_streamed_deltas(self): response = StreamingResponse( [ b'data: {"choices":[{"delta":{"content":"first"}}]}\n', b'data: {"choices":[{"delta":{"content":" second"}}]}\n', b'data: {"choices":[{"message":{"content":"first second"}}]}\n', b"data: [DONE]\n", ] ) with patch("assistant_agent.urllib.request.urlopen", return_value=response): chunks = list( stream_review_assistant( {}, "question", [], "key", "https://example.test/v1", "model" ) ) self.assertEqual(chunks, ["first", " second"]) def test_ndjson_event_writer_flushes_complete_line(self): handler = RequestHandler.__new__(RequestHandler) handler.wfile = io.BytesIO() RequestHandler._write_stream_event(handler, {"type": "delta", "content": "片段"}) self.assertEqual( handler.wfile.getvalue().decode("utf-8"), '{"type":"delta","content":"片段"}\n', ) class ReviewAssistantHistoryTests(unittest.TestCase): def setUp(self) -> None: self.temp = tempfile.TemporaryDirectory() self.database = ReviewDatabase(Path(self.temp.name) / "review.db") self.owner = self.database.create_user("assistant_owner", "salt", "hash") self.other = self.database.create_user("assistant_other", "salt", "hash") def tearDown(self) -> None: self.temp.cleanup() def test_history_is_private_and_clear_is_scoped(self): self.database.save_assistant_exchange( self.owner["id"], "今天怎么看?", "先看承接。", "20260722" ) owner_messages = self.database.list_assistant_messages(self.owner["id"]) self.assertEqual([item["role"] for item in owner_messages], ["user", "assistant"]) self.assertEqual(self.database.list_assistant_messages(self.other["id"]), []) self.assertEqual(self.database.delete_assistant_messages(self.other["id"]), 0) self.assertEqual(self.database.delete_assistant_messages(self.owner["id"]), 2) if __name__ == "__main__": unittest.main()