b005fb1ef6
Some checks are pending
release / release (push) Waiting to run
Package layout (zubble/), vendored benchlib.stream_chat, LocalEnv, minimal CLI, 38 pure-logic tests, and a tag-triggered Forgejo Actions release pipeline that builds a single-file PyInstaller executable.
142 lines
No EOL
4.6 KiB
Python
142 lines
No EOL
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Tests for stream_chat watcher abort hook.
|
|
|
|
The watcher lets a caller abort generation mid-stream when an in-stream
|
|
repetition detector fires. These tests stub only HTTP I/O (the SSE
|
|
response), never the logic that processes deltas or manages the abort.
|
|
|
|
python3 -m unittest discover -s tests -v
|
|
"""
|
|
import os
|
|
import sys
|
|
import json
|
|
import unittest
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from zubble import benchlib # noqa: E402
|
|
from zubble.doom import StreamRepetitionDetector, Level # noqa: E402
|
|
|
|
|
|
class _SSEStub:
|
|
"""A minimal file-like that yields SSE data frames then [DONE]."""
|
|
|
|
def __init__(self, frames):
|
|
self._data = b""
|
|
for payload in frames:
|
|
self._data += f"data: {json.dumps(payload)}\n\n".encode()
|
|
self._data += b"data: [DONE]\n\n"
|
|
self._pos = 0
|
|
|
|
def read(self, n=4096):
|
|
chunk = self._data[self._pos:self._pos + n]
|
|
self._pos += len(chunk)
|
|
return chunk
|
|
|
|
def __iter__(self):
|
|
for line in self._data.split(b"\n"):
|
|
yield line + b"\n"
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *a):
|
|
return False
|
|
|
|
|
|
def _content_frames(text, n, finish="stop", total=None):
|
|
"""Build n content-delta frames + finish + usage."""
|
|
frames = []
|
|
for _ in range(n):
|
|
frames.append({"choices": [{"delta": {"content": text}}]})
|
|
frames.append({"choices": [{"finish_reason": finish}]})
|
|
frames.append({"usage": {"completion_tokens": total or n,
|
|
"prompt_tokens": 10}})
|
|
return frames
|
|
|
|
|
|
class StreamChatWatcherTest(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self._urlopen = benchlib.urllib.request.urlopen
|
|
os.environ["BENCH_MODEL"] = "test-model"
|
|
|
|
def tearDown(self):
|
|
benchlib.urllib.request.urlopen = self._urlopen
|
|
os.environ.pop("BENCH_MODEL", None)
|
|
|
|
def _stub_urlopen(self, frames):
|
|
def _urlopen(req, timeout=2400):
|
|
return _SSEStub(frames)
|
|
benchlib.urllib.request.urlopen = _urlopen
|
|
|
|
def test_watcher_aborts_mid_stream(self):
|
|
self._stub_urlopen(_content_frames("AB ", 100))
|
|
count = [0]
|
|
|
|
def watcher(kind, text):
|
|
count[0] += 1
|
|
return count[0] >= 30
|
|
|
|
content, reasoning, perf = benchlib.stream_chat(
|
|
"http://x", "prompt", timeout=10, watcher=watcher)
|
|
self.assertEqual(perf["finish"], "aborted")
|
|
self.assertEqual(count[0], 30)
|
|
self.assertEqual(len(content), 30 * 3)
|
|
|
|
def test_no_watcher_no_abort(self):
|
|
self._stub_urlopen(_content_frames("hello", 5, total=7))
|
|
content, reasoning, perf = benchlib.stream_chat(
|
|
"http://x", "prompt", timeout=10)
|
|
self.assertNotEqual(perf["finish"], "aborted")
|
|
self.assertEqual(content, "hello" * 5)
|
|
|
|
def test_watcher_receives_reasoning_kind(self):
|
|
frames = []
|
|
for _ in range(50):
|
|
frames.append({"choices": [{"delta": {"reasoning_content": "x "}}]})
|
|
for _ in range(50):
|
|
frames.append({"choices": [{"delta": {"content": "y "}}]})
|
|
frames.append({"choices": [{"finish_reason": "stop"}]})
|
|
frames.append({"usage": {"completion_tokens": 100, "prompt_tokens": 10}})
|
|
self._stub_urlopen(frames)
|
|
|
|
kinds = []
|
|
|
|
def watcher(kind, text):
|
|
kinds.append(kind)
|
|
return len(kinds) >= 60
|
|
|
|
content, reasoning, perf = benchlib.stream_chat(
|
|
"http://x", "prompt", timeout=10, watcher=watcher)
|
|
self.assertEqual(perf["finish"], "aborted")
|
|
self.assertEqual(kinds[:50], ["reasoning"] * 50)
|
|
self.assertIn("content", kinds)
|
|
|
|
def test_watcher_false_does_not_abort(self):
|
|
self._stub_urlopen(_content_frames("X", 10, total=12))
|
|
|
|
def watcher(kind, text):
|
|
return False
|
|
|
|
content, reasoning, perf = benchlib.stream_chat(
|
|
"http://x", "prompt", timeout=10, watcher=watcher)
|
|
self.assertNotEqual(perf["finish"], "aborted")
|
|
self.assertEqual(content, "X" * 10)
|
|
|
|
def test_integration_with_stream_repetition_detector(self):
|
|
detector = StreamRepetitionDetector()
|
|
|
|
def watcher(kind, text):
|
|
result = detector.feed(kind, text)
|
|
return result != Level.NONE
|
|
|
|
self._stub_urlopen(_content_frames("x ", 51 * 50))
|
|
content, reasoning, perf = benchlib.stream_chat(
|
|
"http://x", "prompt", timeout=10, watcher=watcher)
|
|
self.assertEqual(perf["finish"], "aborted")
|
|
self.assertEqual(detector.level, Level.L1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main() |