98 lines
3.7 KiB
Python
98 lines
3.7 KiB
Python
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()
|