business-assistant/scripts/tests/test_deepseek_proxy.py

import contextlib
import importlib.util
import io
import json
import os
import threading
import tempfile
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from unittest.mock import patch


spec = importlib.util.spec_from_file_location(
    "deepseek_proxy_probe", Path(__file__).resolve().parents[1] / "test_deepseek_proxy.py"
)
probe_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(probe_module)


class DeepSeekProxyTests(unittest.TestCase):
    def setUp(self):
        self.requests = []
        self.status = 200
        self.body = json.dumps({"choices": [{"message": {"content": "OK"}}]}).encode()
        self.headers = {"Content-Type": "application/json"}
        test = self

        class Handler(BaseHTTPRequestHandler):
            protocol_version = 'HTTP/1.1'

            def send_body(self):
                self.send_response(test.status)
                for key, value in test.headers.items():
                    self.send_header(key, value)
                if 'Transfer-Encoding' not in test.headers and 'Content-Length' not in test.headers:
                    self.send_header('Content-Length', str(len(test.body)))
                else:
                    # Broken framing cases deliberately terminate the response.
                    self.close_connection = True
                self.end_headers()
                self.wfile.write(test.body)

            def do_GET(self):
                test.requests.append({
                    "path": self.path, "method": "GET",
                    "authorization": self.headers.get("Authorization"),
                    "payload": None, "content_length": self.headers.get("Content-Length"),
                })
                self.send_body()

            def do_POST(self):
                test.requests.append({
                    "path": self.path, "method": "POST",
                    "connection": self.headers.get("Connection"),
                    "authorization": self.headers.get("Authorization"),
                    "payload": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
                })
                self.send_body()

            def log_message(self, *_args):
                pass

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
        self.thread.start()
        self.url = f"http://127.0.0.1:{self.server.server_port}/proxy/?servertype=deepseek&proxyToken=fake%2Btoken%26value"
        self.environment = {
            "DEEPSEEK_BASE_URL": self.url, "DEEPSEEK_API_KEY": "fake-api-key",
            "DEEPSEEK_MODEL": "test-model", "HTTP_PROXY": "http://invalid.example:1",
        }

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=2)

    def run_probe(self, method=None):
        with patch.dict(os.environ, self.environment), patch.object(probe_module, "load_env_values", return_value={
            "DEEPSEEK_BASE_URL": "http://ignored.example/proxy/?serverType=deepseek&proxyToken=ignored",
            "DEEPSEEK_API_KEY": "ignored", "DEEPSEEK_MODEL": "ignored"
        }) as load:
            result = probe_module.probe(timeout=2, **({"method": method} if method else {}))
            self.assertEqual(load.call_args.args[0], Path(__file__).resolve().parents[2] / ".env")
        return result

    def test_actual_http_post_preserves_fixed_url_and_payload(self):
        result = self.run_probe()
        self.assertTrue(result["ok"], result)
        request = self.requests[0]
        self.assertEqual(request["path"], "/proxy/?servertype=deepseek&proxyToken=fake%2Btoken%26value",
                         "固定代理入口不可追加/chat/completions或修改token")
        self.assertEqual(request["authorization"], "Bearer fake-api-key")
        self.assertEqual(request["payload"]["model"], "test-model")
        self.assertEqual(request["payload"]["messages"], [{"role": "user", "content": "1+1等于多少"}])
        self.assertEqual(request["connection"], "keep-alive")
        self.assertEqual(result["transport"], "http.client")
        self.assertFalse(request["payload"]["stream"])
        self.assertTrue(result["environment_proxy_configured"])
        self.assertNotIn("fake+token&value", json.dumps(result))
        self.assertEqual(len(self.requests), 1)

    def test_get_sends_exact_gateway_url_without_chat_body(self):
        result = self.run_probe("GET")
        self.assertEqual(result["method"], "GET")
        self.assertFalse(result["chat_payload_sent"])
        self.assertTrue(result["http_ok"])
        self.assertEqual(self.requests[0]["path"],
                         "/proxy/?servertype=deepseek&proxyToken=fake%2Btoken%26value")
        self.assertEqual(self.requests[0]["method"], "GET")
        self.assertIsNone(self.requests[0]["content_length"])
        self.assertIsNone(self.requests[0]["payload"])
        self.assertEqual(len(self.requests), 1)

    def test_complete_chat_path_removes_trailing_slash_without_changing_query(self):
        self.environment['DEEPSEEK_BASE_URL'] = self.url.replace('/proxy/?', '/proxy/chat/completions/?')
        result = self.run_probe()
        self.assertTrue(result['ok'])
        self.assertEqual(self.requests[0]['path'],
                         '/proxy/chat/completions?servertype=deepseek&proxyToken=fake%2Btoken%26value')
        self.assertTrue(result['endpoint'].endswith('/proxy/chat/completions'))

    def test_get_http_success_is_not_enough_to_verify_chat(self):
        self.body = b"gateway connected"
        result = self.run_probe("GET")
        self.assertTrue(result["http_ok"])
        self.assertFalse(result["ok"])
        self.assertEqual(result["code"], "unexpected_response")

    def test_gateway_auth_failure_redacts_echoed_credentials(self):
        self.status = 401
        self.body = b"Authentication Fails (governor) fake-api-key fake+token&value fake%2Btoken%26value"
        result = self.run_probe()
        self.assertEqual(result["code"], "proxy_authentication_failed")
        self.assertEqual(result["http_status"], 401)
        for secret in ("fake-api-key", "fake+token&value", "fake%2Btoken%26value"):
            self.assertNotIn(secret, json.dumps(result))
        self.assertEqual(len(self.requests), 1, "诊断只发一次请求，不自动重试")

    def test_http_200_non_chat_response_is_failure(self):
        for body in (b"gateway reachable", b'{"error":{"message":"denied"}}', b'{"choices":[]}'):
            with self.subTest(body=body):
                self.body = body
                result = self.run_probe()
                self.assertFalse(result["ok"])
                self.assertEqual(result["code"], "unexpected_response")

    def set_unfinished_chunked_body(self, body):
        self.headers["Transfer-Encoding"] = "chunked"
        # A complete data chunk followed by a closed connection, without the
        # mandatory zero-sized final chunk. This reproduces the gateway symptom.
        self.body = f"{len(body):X}\r\n".encode() + body + b"\r\n"

    def test_http_500_incomplete_read_preserves_redacted_error_body(self):
        self.status = 500
        error = b'{"error":"upstream failed fake-api-key fake+token&value fake%2Btoken%26value"}'
        self.set_unfinished_chunked_body(error)
        result = self.run_probe()
        self.assertEqual(result["http_status"], 500)
        self.assertEqual(result["code"], "http_error")
        self.assertEqual(result["error_type"], "IncompleteRead")
        self.assertFalse(result["response_complete"])
        self.assertEqual(result["received_bytes"], len(error))
        self.assertIn("upstream failed", result["response_preview"])
        for secret in ("fake-api-key", "fake+token&value", "fake%2Btoken%26value"):
            self.assertNotIn(secret, json.dumps(result))
        self.assertEqual(len(self.requests), 1)

    def test_incomplete_http_200_valid_json_cannot_report_chat_success(self):
        self.set_unfinished_chunked_body(self.body)
        result = self.run_probe()
        self.assertFalse(result["ok"], "响应体JSON合法也不能把不完整HTTP响应报成功")
        self.assertEqual(result["code"], "response_incomplete")
        self.assertIn("choices", result["response_preview"])

    def test_short_content_length_error_is_visible(self):
        self.status = 500
        self.body = b'{"error":"short body"}'
        self.headers["Content-Length"] = str(len(self.body) + 20)
        result = self.run_probe()
        self.assertEqual(result["code"], "http_error")
        self.assertEqual(result["error_type"], "ContentLengthMismatch")
        self.assertFalse(result["response_complete"])
        self.assertIn("short body", result["response_preview"])
        self.assertEqual(result["content_length"] - result["received_bytes"], 20)

    def test_redirect_is_reported_without_second_request(self):
        self.status = 302
        self.headers["Location"] = f"http://127.0.0.1:{self.server.server_port}/other"
        result = self.run_probe()
        self.assertEqual(result["code"], "http_error")
        self.assertEqual(len(self.requests), 1)

    def test_missing_api_key_can_probe_gateway_without_fake_bearer(self):
        self.environment["DEEPSEEK_API_KEY"] = ""
        result = self.run_probe()
        self.assertFalse(result["api_key_configured"])
        self.assertIsNone(self.requests[0]["authorization"])

    def test_invalid_and_wrong_case_params_do_not_send(self):
        for url, code in ((self.url.replace("servertype", "serverType"), "missing_proxy_params"),
                          (self.url.replace("proxyToken", "proxytoken"), "missing_proxy_params"),
                          ("http://gateway.example\n:123/proxy", "invalid_base_url"),
                          ("http://gateway.example:bad/proxy", "invalid_base_url")):
            with self.subTest(code=code, url=url):
                self.environment["DEEPSEEK_BASE_URL"] = url
                self.assertEqual(self.run_probe()["code"], code)
        self.assertEqual(self.requests, [])

    def test_connection_failure_does_not_echo_full_url(self):
        with patch.object(probe_module, "HTTPConnection") as build:
            build.return_value.request.side_effect = ConnectionRefusedError(self.url)
            result = self.run_probe()
            build.return_value.close.assert_called_once()
        self.assertEqual(result["code"], "connection_failed")
        self.assertNotIn("fake", json.dumps(result))

    def test_https_timeout_closes_connection_without_secret_echo(self):
        self.environment['DEEPSEEK_BASE_URL'] = self.url.replace('http:', 'https:')
        with patch.object(probe_module, 'HTTPSConnection') as build:
            build.return_value.getresponse.side_effect = TimeoutError(self.url)
            result = self.run_probe()
            self.assertEqual(result['code'], 'timeout')
            build.return_value.close.assert_called_once()
            self.assertEqual(build.call_args.kwargs['timeout'], 2)
        self.assertNotIn('fake', json.dumps(result))

    def test_complete_chunked_response(self):
        self.set_unfinished_chunked_body(self.body)
        self.body += b'0\r\n\r\n'
        result = self.run_probe()
        self.assertTrue(result['ok'], result)
        self.assertTrue(result['response_complete'])

    def test_cli_exit_codes_and_safe_json(self):
        for ok in (True, False):
            output = io.StringIO()
            with patch.object(probe_module, "probe", return_value={"ok": ok}), patch.object(
                probe_module.sys, "argv", ["test_deepseek_proxy.py"]
            ), contextlib.redirect_stdout(output):
                self.assertEqual(probe_module.main(), 0 if ok else 1)
            self.assertEqual(json.loads(output.getvalue()), {"ok": ok})

    def test_cli_defaults_to_post_and_method_can_be_selected(self):
        for arguments, method in (([], "POST"), (["--method", "GET"], "GET"),
                                  (["--method", "POST"], "POST")):
            with patch.object(probe_module, "probe", return_value={"ok": True}) as probe, patch.object(
                probe_module.sys, "argv", ["test_deepseek_proxy.py", *arguments]
            ), contextlib.redirect_stdout(io.StringIO()):
                self.assertEqual(probe_module.main(), 0)
                self.assertEqual(probe.call_args.kwargs["method"], method)


class EnvReadTests(unittest.TestCase):
    def test_synthetic_env_preserves_case_quoting_and_last_definition(self):
        with tempfile.TemporaryDirectory() as temporary:
            env_path = Path(temporary) / ".env"
            env_path.write_text(
                '# synthetic settings\nDEEPSEEK_BASE_URL=http://old.example\n'
                'export DEEPSEEK_BASE_URL="http://gateway.example/proxy/?servertype=deepseek&proxyToken=fake-token" # comment\n'
                "DEEPSEEK_API_KEY='fake-key'\nDEEPSEEK_MODEL=deepseek-flash # comment\nUNRELATED=ignored\n",
                encoding="utf-8-sig",
            )
            values = probe_module.load_env_values(env_path)
            self.assertEqual(values["DEEPSEEK_BASE_URL"],
                             "http://gateway.example/proxy/?servertype=deepseek&proxyToken=fake-token")
            self.assertEqual(values["DEEPSEEK_API_KEY"], "fake-key")
            self.assertEqual(values["DEEPSEEK_MODEL"], "deepseek-flash")
            self.assertEqual(len(values), 3)
            env_path.write_text('DEEPSEEK_BASE_URL="unclosed', encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "invalid_env_quoting"):
                probe_module.load_env_values(env_path)


if __name__ == "__main__":
    unittest.main()