109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from mentor_agent import MentorSkill, chat_with_mentor, stream_with_mentor
|
|
|
|
|
|
class FakeStreamResponse:
|
|
def __init__(self, lines: list[bytes]) -> None:
|
|
self.lines = lines
|
|
|
|
def __enter__(self):
|
|
return iter(self.lines)
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
return False
|
|
|
|
|
|
class MentorStreamTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.skill = MentorSkill(
|
|
skill_id="test-mentor",
|
|
name="测试老师",
|
|
description="测试",
|
|
tagline="先看事实",
|
|
focus=("纪律",),
|
|
content="只做条件化判断。",
|
|
path=Path("SKILL.md"),
|
|
)
|
|
self.lines = [
|
|
b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
|
|
b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
|
|
b"data: [DONE]\n",
|
|
]
|
|
|
|
def test_stream_requests_upstream_streaming_and_yields_deltas(self):
|
|
captured = {}
|
|
|
|
def open_request(request, timeout):
|
|
captured["payload"] = json.loads(request.data.decode("utf-8"))
|
|
captured["accept"] = request.headers.get("Accept")
|
|
return FakeStreamResponse(self.lines)
|
|
|
|
with patch("mentor_agent.urllib.request.urlopen", side_effect=open_request):
|
|
chunks = list(
|
|
stream_with_mentor(
|
|
self.skill, {"data_trade_date": "20260723"}, "怎么看？", [],
|
|
"key", "https://example.test/v1", "model",
|
|
)
|
|
)
|
|
|
|
self.assertEqual(chunks, ["first", " second"])
|
|
self.assertTrue(captured["payload"]["stream"])
|
|
self.assertEqual(captured["accept"], "text/event-stream")
|
|
|
|
def test_non_streaming_compatibility_wrapper_collects_chunks(self):
|
|
with patch(
|
|
"mentor_agent.urllib.request.urlopen",
|
|
return_value=FakeStreamResponse(self.lines),
|
|
):
|
|
result = chat_with_mentor(
|
|
self.skill, {}, "怎么看？", [], "key", "https://example.test/v1", "model"
|
|
)
|
|
self.assertEqual(result["answer"], "first second")
|
|
|
|
def test_final_full_message_does_not_duplicate_streamed_deltas(self):
|
|
lines = [
|
|
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(
|
|
"mentor_agent.urllib.request.urlopen",
|
|
return_value=FakeStreamResponse(lines),
|
|
):
|
|
chunks = list(
|
|
stream_with_mentor(
|
|
self.skill, {}, "question", [], "key",
|
|
"https://example.test/v1", "model",
|
|
)
|
|
)
|
|
self.assertEqual(chunks, ["first", " second"])
|
|
|
|
def test_snapshot_only_stream_emits_only_new_suffix(self):
|
|
lines = [
|
|
b'data: {"choices":[{"message":{"content":"first"}}]}\n',
|
|
b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
|
|
b"data: [DONE]\n",
|
|
]
|
|
with patch(
|
|
"mentor_agent.urllib.request.urlopen",
|
|
return_value=FakeStreamResponse(lines),
|
|
):
|
|
chunks = list(
|
|
stream_with_mentor(
|
|
self.skill, {}, "question", [], "key",
|
|
"https://example.test/v1", "model",
|
|
)
|
|
)
|
|
self.assertEqual(chunks, ["first", " second"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|