tests/test_streaming.py

"""流式与思考链路的端到端计时验证。

起一个真实的 uvicorn，用一个「慢慢吐分片」的假 Gemini 后端，然后用 httpx 真流式
接收，测量首字节时间（TTFB）与总耗时。如果 TTFB 远小于总耗时，说明服务端确实在
逐块下发；如果两者接近，说明中间被整体缓冲了。

同时验证 ENABLE_THINKING=true 时 thoughts_delta 会以 reasoning_content 下发。
"""

import asyncio
import os
import socket
import sys
import tempfile
import threading
import time
from pathlib import Path
from types import SimpleNamespace

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

os.environ["SECURE_1PSID"] = "PSID-STREAM"
os.environ["SECURE_1PSIDTS"] = "TS-STREAM"
os.environ["GEMINI_COOKIE_PATH"] = tempfile.mkdtemp(prefix="gs-stream-")
os.environ["ENABLE_THINKING"] = "true"
os.environ["SESSION_HEALTH_INTERVAL"] = "3600"

import httpx  # noqa: E402

import gemini_session as gs  # noqa: E402

PASS, FAIL = [], []


def check(name, cond, extra=""):
	(PASS if cond else FAIL).append(name)
	print(("[OK]   " if cond else "[FAIL] ") + name + (f"  {extra}" if extra and not cond else ""))


CHUNK_DELAY = 0.35
THOUGHTS = ["先理解问题。", "再给出结论。"]
TEXTS = ["Hello ", "world ", "from ", "streaming ", "test "]


class FakeCookies:
	def __init__(self, values):
		self.jar = [SimpleNamespace(name=k, value=v) for k, v in values.items()]

	def get(self, name, default=None):
		for c in self.jar:
			if c.name == name:
				return c.value
		return default


class FakeChunk:
	def __init__(self, text_delta=None, thoughts_delta=None):
		self.text_delta = text_delta
		self.thoughts_delta = thoughts_delta
		self.text = text_delta or ""
		self.thoughts = None
		self.images = []
		self.metadata = ["cid-stream"]


class FakeResponse:
	def __init__(self):
		self.text = "".join(TEXTS)
		self.thoughts = "".join(THOUGHTS)
		self.images = []
		self.metadata = ["cid-generate"]


class FakeClient:
	def __init__(self, psid=None, psidts=None, **kwargs):
		self.psid = psid
		self.psidts = psidts
		self._running = False
		self.client = None
		self.auto_refresh = False
		self.refresh_task = None
		self.activity_task = None
		self.account_status = SimpleNamespace(name="AVAILABLE")
		self.cookies = FakeCookies({"__Secure-1PSID": psid, "__Secure-1PSIDTS": psidts})

	async def init(self, **kwargs):
		self._running = True
		self.client = object()
		self.auto_refresh = kwargs.get("auto_refresh", False)
		if self.auto_refresh:
			self.refresh_task = asyncio.create_task(self.start_auto_refresh())
			self.activity_task = asyncio.create_task(self.start_activity_watchdog())

	async def generate_content(self, prompt, **kwargs):
		return FakeResponse()

	async def generate_content_stream(self, prompt, **kwargs):
		for thought in THOUGHTS:
			await asyncio.sleep(CHUNK_DELAY)
			yield FakeChunk(thoughts_delta=thought)
		for text in TEXTS:
			await asyncio.sleep(CHUNK_DELAY)
			yield FakeChunk(text_delta=text)

	async def fetch_latest_chat_response(self, cid):
		return FakeResponse()

	async def delete_chat(self, cid):
		return None

	async def start_auto_refresh(self):
		while self._running:
			await asyncio.sleep(3600)

	async def start_activity_watchdog(self):
		while self._running:
			await asyncio.sleep(3600)

	async def close(self):
		self._running = False
		self.client = None


def free_port() -> int:
	with socket.socket() as s:
		s.bind(("127.0.0.1", 0))
		return s.getsockname()[1]


def main():
	gs.GeminiClient = FakeClient

	import uvicorn

	import main as app_main

	port = free_port()
	config = uvicorn.Config(app_main.app, host="127.0.0.1", port=port, log_level="warning")
	server = uvicorn.Server(config)
	threading.Thread(target=server.run, daemon=True).start()

	base = f"http://127.0.0.1:{port}"
	deadline = time.time() + 20
	while time.time() < deadline:
		try:
			if httpx.get(f"{base}/", timeout=1).status_code == 200:
				break
		except Exception:
			time.sleep(0.1)
	else:
		print("[FAIL] 服务未能启动")
		return 1

	check("1 uvicorn 已就绪", True)

	payload = {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}], "stream": True}

	with httpx.Client(timeout=60) as client:
		start = time.time()
		arrivals = []
		reasoning_seen = []
		all_lines = []
		with client.stream("POST", f"{base}/v1/chat/completions", json=payload) as response:
			check("2 HTTP 200", response.status_code == 200, str(response.status_code))
			check("3 Content-Type 为 text/event-stream", "text/event-stream" in response.headers.get("content-type", ""),
			      response.headers.get("content-type", ""))
			check("4 未整体缓冲（无 Content-Length）", "content-length" not in {k.lower() for k in response.headers},
			      "响应带 Content-Length 说明被整体缓冲了")
			for line in response.iter_lines():
				if not line.startswith("data: "):
					continue
				arrivals.append(time.time() - start)
				all_lines.append(line)
				if "reasoning_content" in line:
					reasoning_seen.append(line)

		ttfb = arrivals[0] if arrivals else -1
		total = arrivals[-1] if arrivals else -1
		spread = total - ttfb

	print(f"       分片数={len(arrivals)}  TTFB={ttfb:.2f}s  末片={total:.2f}s  跨度={spread:.2f}s")

	check("5 收到了多个 SSE 分片", len(arrivals) >= 5, f"只收到 {len(arrivals)} 个")
	check("6 首字节足够早（真流式）", 0 < ttfb < 1.5, f"TTFB={ttfb:.2f}s")
	check("7 分片随时间陆续到达（非一次性下发）", spread > 1.0, f"跨度仅 {spread:.2f}s")
	check("8 思考内容以 reasoning_content 下发", bool(reasoning_seen), f"匹配 {len(reasoning_seen)} 条")
	check("9 思考段落以 <think> 标记开头", any("<think>" in line for line in all_lines))
	check("9b 思考段落以 </think> 收尾", any("</think>" in line for line in all_lines))

	# 非流式对照：思考应写进 content 的 <think> 块
	with httpx.Client(timeout=30) as client:
		r = client.post(f"{base}/v1/chat/completions",
		                json={**payload, "stream": False})
		body = r.json()
		content = body["choices"][0]["message"]["content"]
	check("10 非流式把思考包在 <think> 里", "<think>" in content and "先理解问题" in content, content[:120])

	server.should_exit = True
	time.sleep(0.5)

	print(f"\n通过 {len(PASS)} / {len(PASS) + len(FAIL)}")
	if FAIL:
		print("失败项：")
		for name in FAIL:
			print("  -", name)
	return 1 if FAIL else 0


if __name__ == "__main__":
	sys.exit(main())