business-assistant/scripts/tests/test_deepseek_proxy_stream.py

import importlib.util
import json
import os
from pathlib import Path
import time
import unittest
from unittest.mock import patch

from scripts.tests import test_deepseek_proxy as ordinary

spec = importlib.util.spec_from_file_location('stream_probe', Path(__file__).resolve().parents[1] / 'test_deepseek_proxy_stream.py')
probe = importlib.util.module_from_spec(spec)
spec.loader.exec_module(probe)


def frame(text=None, finish=None):
    return ('data: ' + json.dumps({'choices': [{'index': 0, 'delta': {'content': text},
                'finish_reason': finish}]}, ensure_ascii=False) + '\r\n\r\n').encode()


class StreamProbeTests(unittest.TestCase):
    setUp = ordinary.DeepSeekProxyTests.setUp
    tearDown = ordinary.DeepSeekProxyTests.tearDown

    def run_probe(self):
        with patch.dict(os.environ, self.environment), patch.object(probe, 'load_env_values', return_value={}):
            return probe.probe(timeout=2)

    def test_real_chunked_sse_keepalive(self):
        self.headers = {'Content-Type': 'text/event-stream; charset=utf-8', 'Transfer-Encoding': 'chunked'}
        frames = [b': heartbeat\r\n\r\n', frame('你'), frame('好'), frame(finish='stop'),
                  b'data: {"choices":[],"usage":{}}\r\n\r\n', b'data: [DONE]\r\n\r\n']
        self.body = b''.join(f'{len(x):x}\r\n'.encode() + x + b'\r\n' for x in frames) + b'0\r\n\r\n'
        result = self.run_probe()
        self.assertTrue(result['ok'], result)
        self.assertTrue(result['done_received'])
        self.assertTrue(result['response_complete'])
        self.assertEqual(result['assistant_text'], '你好')
        self.assertEqual(result['content_chunks'], 2)
        self.assertEqual(self.requests[0]['connection'], 'keep-alive')
        self.assertTrue(self.requests[0]['payload']['stream'])
        self.assertEqual(self.requests[0]['path'], '/proxy/?servertype=deepseek&proxyToken=fake%2Btoken%26value')

    def test_failure_matrix(self):
        self.headers = {'Content-Type': 'text/event-stream'}
        for body, code in (
            (frame('x', 'stop'), 'stream_incomplete'),
            (frame('x') + b'data: [DONE]\n\n', 'stream_incomplete'),
            (frame(None, 'stop') + b'data: [DONE]\n\n', 'empty_stream_content'),
            (frame('x', 'length') + b'data: [DONE]\n\n', 'incomplete_completion'),
            (b'data: {broken}\n\n', 'invalid_sse'),
            (b'data: {"error":"fake-api-key"}\n\n', 'provider_error'),
        ):
            with self.subTest(code=code):
                self.body = body
                result = self.run_probe()
                self.assertFalse(result['ok'])
                self.assertEqual(result['code'], code)
                self.assertNotIn('fake-api-key', json.dumps(result))

    def test_done_cannot_hide_broken_http_chunking(self):
        self.headers = {'Content-Type': 'text/event-stream', 'Transfer-Encoding': 'chunked'}
        body = frame('x', 'stop') + b'data: [DONE]\n\n'
        self.body = f'{len(body):x}\r\n'.encode() + body + b'\r\n'
        result = self.run_probe()
        self.assertFalse(result['ok'])
        self.assertEqual(result['code'], 'response_incomplete')

    def test_json_fallback_and_http_error_are_not_stream_success(self):
        self.assertEqual(self.run_probe()['code'], 'unexpected_response')
        self.status = 500
        self.assertEqual(self.run_probe()['code'], 'http_error')

    def test_bytewise_utf8_multiline_and_redaction(self):
        wire = frame('你好 fake-api-key', 'stop') + b'data: [DONE]\n\n'
        class Response:
            def __init__(self):
                self.parts = iter(bytes([b]) for b in wire)
            def read1(self, _size):
                return next(self.parts, b'')
        result = {'ok': False, 'content_length': len(wire), 'transfer_encoding': ''}
        probe.read_sse(Response(), result, ['fake-api-key'], time.perf_counter(), 2)
        self.assertTrue(result['ok'])
        self.assertNotIn('fake-api-key', result['assistant_text'])

    def test_cli_and_timeout_cleanup(self):
        import contextlib
        import io
        with patch.object(probe.sys, 'argv', ['test_deepseek_proxy_stream.py']), patch.object(probe, 'probe', return_value={'ok': True}), contextlib.redirect_stdout(io.StringIO()):
            self.assertEqual(probe.main(), 0)
        with patch.object(probe, 'HTTPConnection') as connection:
            connection.return_value.getresponse.side_effect = TimeoutError
            self.assertEqual(self.run_probe()['code'], 'timeout')
            connection.return_value.close.assert_called_once()