zubble/tests/test_stream_watcher.py

#!/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()