business-assistant/scripts/tests/test_deepseek_transport.py

"""Production clients against a synthetic gateway; no real settings or network."""
import contextlib
import io
import json
import os
from pathlib import Path
import tempfile
import threading
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import patch

from scripts import deploy_check
from src.step1_knowledge_prepare.summarize import Config, chat_endpoint, load_config, transport, RequestFailure
from src.step3_question_answer.streaming import stream_transport


class EndpointTests(unittest.TestCase):
    def test_base_and_complete_urls_preserve_query(self):
        for base, expected in (
            ('https://example.invalid', 'https://example.invalid/chat/completions'),
            ('https://example.invalid/v1/', 'https://example.invalid/v1/chat/completions'),
            ('https://example.invalid/chat/completions', 'https://example.invalid/chat/completions'),
            ('http://example.invalid:123/proxy/', 'http://example.invalid:123/proxy/chat/completions'),
            ('http://example.invalid/proxy/chat/completions/', 'http://example.invalid/proxy/chat/completions'),
        ):
            query = '?servertype=deepseek&proxyToken=fake%2B%26%2f&x=1&x=2'
            with self.subTest(base=base):
                self.assertEqual(chat_endpoint(base + query), expected + query)

    def test_invalid_urls_have_safe_errors(self):
        for url in ('file:///secret', 'http://user:secret@example.invalid',
                    'http://example.invalid:bad', 'http://example.invalid/#secret',
                    'http://example.invalid/\nsecret', 'http://', 'http://example.invalid\\secret'):
            with self.subTest(url=url), self.assertRaisesRegex(ValueError, '^invalid_deepseek_base_url$'):
                chat_endpoint(url)


class GatewayTests(unittest.TestCase):
    def setUp(self):
        self.calls = []
        self.mode = 'normal'
        self.request_headers = []
        self.first_received = threading.Event()
        test = self

        class Handler(BaseHTTPRequestHandler):
            protocol_version = 'HTTP/1.1'
            def log_message(self, *_):
                pass

            def do_POST(self):
                payload = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
                test.calls.append((self.path, self.headers.get('Authorization'), payload))
                test.request_headers.append(dict(self.headers))
                if test.mode == 'redirect':
                    self.send_response(302)
                    self.send_header('Location', '/other')
                    self.send_header('Content-Length', '0')
                    self.end_headers()
                    return
                self.send_response(200)
                if test.mode == 'truncated':
                    self.send_header('Transfer-Encoding', 'chunked')
                    self.end_headers()
                    self.wfile.write(b'broken\r\n')
                    self.close_connection = True
                    return
                if payload.get('stream'):
                    self.send_header('Content-Type', 'text/event-stream')
                    self.send_header('Transfer-Encoding', 'chunked')
                    self.end_headers()
                    def write_chunk(data):
                        self.wfile.write(f'{len(data):x}\r\n'.encode() + data + b'\r\n')
                        self.wfile.flush()
                    if test.mode == 'incremental':
                        write_chunk(b'data: {"choices":[{"delta":{"content":"first"}}]}\n\n')
                        if not test.first_received.wait(2):
                            self.close_connection = True
                            return
                    frame = {'choices': [{'delta': {'content': 'OK'}, 'finish_reason': 'stop'}]}
                    write_chunk(('data: ' + json.dumps(frame) + '\n\ndata: [DONE]\n\n').encode())
                    if test.mode == 'done_truncated':
                        self.close_connection = True
                    else:
                        self.wfile.write(b'0\r\n\r\n')
                else:
                    body = b'{"choices":[{"message":{"content":"OK"}}]}'
                    self.send_header('Content-Length', str(len(body)))
                    self.end_headers()
                    self.wfile.write(body)

        self.server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.temp = tempfile.TemporaryDirectory()
        self.root = Path(self.temp.name)
        self.url = (f'http://127.0.0.1:{self.server.server_port}/proxy/chat/completions'
                    '?servertype=deepseek&proxyToken=fake%2B%26value')
        (self.root / '.env').write_text('DEEPSEEK_BASE_URL=' + self.url +
            '\nDEEPSEEK_API_KEY=fake-key\nDEEPSEEK_MODEL=fake-model\n', encoding='utf-8')
        self.env = patch.dict(os.environ, {'HTTP_PROXY': 'http://127.0.0.1:1',
            'http_proxy': 'http://127.0.0.1:1', 'NO_PROXY': '', 'no_proxy': ''}, clear=True)
        self.env.start()

    def tearDown(self):
        self.env.stop()
        self.server.shutdown()
        self.server.server_close()
        self.thread.join()
        self.temp.cleanup()

    def test_production_json_stream_and_deployment_check(self):
        config = load_config(self.root / '.env')
        payload = {'model': config.model, 'messages': [{'role': 'user', 'content': 'synthetic'}]}
        self.assertEqual(transport(payload, config)['choices'][0]['message']['content'], 'OK')
        parts = []
        stream_transport(payload, config, parts.append)
        self.assertEqual(parts, ['OK'])
        output = io.StringIO()
        with contextlib.redirect_stdout(output):
            self.assertEqual(deploy_check.run(self.root, action='check-deepseek'), 0)
        self.assertTrue(json.loads(output.getvalue())['checks']['deepseek']['ok'])
        self.assertEqual(len(self.calls), 3)
        for headers in self.request_headers:
            self.assertEqual(headers['Connection'], 'keep-alive')
            self.assertEqual(headers['Accept'], '*/*')
        for path, auth, body in self.calls:
            self.assertEqual(path, '/proxy/chat/completions?servertype=deepseek&proxyToken=fake%2B%26value')
            self.assertEqual(auth, 'Bearer fake-key')
            self.assertEqual(body['model'], 'fake-model')
            self.assertTrue(body['messages'])
        self.assertNotIn('fake-key', repr(config))
        self.assertNotIn('proxyToken', output.getvalue())

    def test_truncated_response_and_redirect_report_safe_codes(self):
        for mode, code in (('truncated', 'response_incomplete'), ('redirect', 'http_302')):
            self.mode = mode
            output = io.StringIO()
            with contextlib.redirect_stdout(output):
                self.assertEqual(deploy_check.run(self.root, action='check-deepseek'), 1)
            self.assertEqual(json.loads(output.getvalue())['checks']['deepseek']['code'], code)
            stream_code = code
            with self.assertRaisesRegex(RequestFailure, '^' + stream_code + '$'):
                stream_transport({}, load_config(self.root / '.env'), lambda _: None)
        self.assertEqual(len(self.calls), 4)

    def test_first_delta_is_delivered_before_server_finishes(self):
        self.mode = 'incremental'
        parts = []
        def consume(text):
            parts.append(text)
            self.first_received.set()
        response = stream_transport({}, load_config(self.root / '.env'), consume)
        self.assertEqual(parts, ['first', 'OK'])
        self.assertEqual(response['choices'][0]['message']['content'], 'firstOK')

    def test_done_does_not_mask_broken_chunked_end(self):
        self.mode = 'done_truncated'
        with self.assertRaisesRegex(RequestFailure, '^response_incomplete$'):
            stream_transport({}, load_config(self.root / '.env'), lambda _: None)

    def test_gateway_connections_closed_on_timeout_and_callback_failure(self):
        from src.step1_knowledge_prepare import summarize
        config = load_config(self.root / '.env')
        with patch.object(summarize, 'HTTPConnection') as factory:
            factory.return_value.getresponse.side_effect = TimeoutError('secret-url')
            with self.assertRaisesRegex(RequestFailure, '^request_timeout$'):
                transport({}, config)
            factory.return_value.close.assert_called_once()
        with patch.object(summarize, 'HTTPConnection') as factory:
            response = factory.return_value.getresponse.return_value
            response.__enter__.return_value = response
            response.status = 200
            response.read1.return_value = b'data: {"choices":[{"delta":{"content":"x"}}]}\n\n'
            def cancel(_):
                raise RequestFailure('stream_cancelled', False)
            with self.assertRaisesRegex(RequestFailure, '^stream_cancelled$'):
                stream_transport({}, config, cancel)
            factory.return_value.close.assert_called_once()