tests/test_integration.py

"""端到端集成验证：真实 FastAPI 应用 + 假 Gemini 后端，不联网。"""

import asyncio
import os
import sys
import tempfile
import time
from types import SimpleNamespace

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

TMP = tempfile.mkdtemp(prefix="gs-e2e-")
os.environ["SECURE_1PSID"] = "PSID-E2E"
os.environ["SECURE_1PSIDTS"] = "ENV-TS-E2E"
os.environ["GEMINI_COOKIE_PATH"] = TMP
os.environ["SESSION_HEALTH_INTERVAL"] = "30"
os.environ["AUTO_DELETE_CHAT"] = "true"

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 ""))


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 FakeResponse:
	def __init__(self, text, cid, delta=None):
		self.text = text
		self.text_delta = delta
		self.thoughts = None
		self.images = []
		self.metadata = [cid]


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})
		self.deleted = []
		self.fail_next_generate = False

	async def init(self, **kwargs):
		self._running = True
		self.client = object()
		self.auto_refresh = kwargs.get("auto_refresh", False)
		# 对齐真实 GeminiClient.init()：开启 auto_refresh 时立即拉起后台任务
		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):
		if self.fail_next_generate:
			self.fail_next_generate = False
			self.account_status = SimpleNamespace(name="UNAUTHENTICATED")
			raise gs_gemini_error()
		return FakeResponse("OK", "cid-generate")

	async def generate_content_stream(self, prompt, **kwargs):
		yield FakeResponse("Hel", "cid-stream-1", delta="Hel")
		yield FakeResponse("lo", "cid-stream-2", delta="lo")

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

	async def delete_chat(self, cid):
		self.deleted.append(cid)

	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._closed_marker = True
		self._running = False
		self.client = None


def gs_gemini_error():
	from gemini_webapi import GeminiError

	return GeminiError("Failed to generate contents. No output data found in response.")


def main():
	gs.GeminiClient = FakeClient

	import main as app_main  # 导入时才读取环境变量
	from fastapi.testclient import TestClient

	with TestClient(app_main.app) as client:
		# ---------------------------------------------------------- 生命周期
		check("E1 lifespan 启动了监督协程", app_main.session_manager.status()["supervisor_alive"] is True)

		# 初始化是后台任务，等它跑完再断言
		deadline = time.time() + 10
		while time.time() < deadline and not app_main.session_manager.status()["initialized"]:
			time.sleep(0.1)
		check("E1b 后台初始化已完成", app_main.session_manager.status()["initialized"] is True)

		# ---------------------------------------------------------- /api/session
		r = client.get("/api/session")
		check("E2 /api/session 返回 200", r.status_code == 200, str(r.status_code))
		data = r.json()
		check("E3 基于 Cookie 初始化成功", data["initialized"] is True)
		check("E4 凭据来源为 environment", data["credential_source"] == "environment")
		check("E5 轮换任务存活", data["rotation_task_alive"] is True)
		check("E6 1PSIDTS 已脱敏", "*" in data["live_1psidts_masked"] and data["live_1psidts_masked"] != "ENV-TS-E2E")
		check("E7 缓存已落盘", data["cached_1psidts_masked"] != "")

		# ---------------------------------------------------------- 非流式对话
		payload = {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}]}
		r = client.post("/v1/chat/completions", json=payload)
		check("E8 非流式请求 200", r.status_code == 200, f"{r.status_code} {r.text[:200]}")
		body = r.json()
		check("E9 返回内容正确", body["choices"][0]["message"]["content"].strip() == "OK", str(body)[:200])

		# ---------------------------------------------------------- 流式对话
		r = client.post("/v1/chat/completions", json={**payload, "stream": True})
		check("E10 流式请求 200", r.status_code == 200, str(r.status_code))
		check("E11 流式返回 SSE 数据", "data: [DONE]" in r.text and "Hel" in r.text, r.text[:200])

		# ---------------------------------------------------------- 失效自愈（非流式）
		bad = app_main.session_manager.client
		bad.fail_next_generate = True
		r = client.post("/v1/chat/completions", json=payload)
		check("E12 会话失效后请求仍成功（自动重建重试）", r.status_code == 200, f"{r.status_code} {r.text[:200]}")
		check("E13 确实重建了客户端", app_main.session_manager.client is not bad)
		check("E14 recycle_count 递增", app_main.session_manager.status()["recycle_count"] >= 1)

		# ---------------------------------------------------------- 轮换 -> 缓存 -> 状态
		live = app_main.session_manager.client
		live.cookies = FakeCookies({"__Secure-1PSID": "PSID-E2E", "__Secure-1PSIDTS": "ROTATED-E2E"})
		app_main.session_manager._observe_rotation(live)
		r = client.get("/api/session")
		data = r.json()
		check("E15 轮换被记录", data["rotation_count"] >= 1)
		check("E16 轮换值已落盘", data["cached_1psidts_masked"].startswith("ROTA"))
		check("E17 换值年龄已更新", data["last_1psidts_change_age_sec"] is not None and data["last_1psidts_change_age_sec"] < 60)

		# ---------------------------------------------------------- 管理面板状态
		# 没有 API_KEY 时管理面板不可用，这里直接调用函数绕过 token
		import admin

		status = asyncio.run(admin.get_status(token=""))
		check("E18 管理面板状态含 session 块", "session" in status and "rotation_count" in status)
		check("E19 管理面板透出实时 1PSIDTS", status["live_1psidts_masked"] != "")

		# ---------------------------------------------------------- 关闭
	check("E20 退出 lifespan 后监督协程停止", app_main.session_manager.status()["supervisor_alive"] is False)

	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())