tests/test_models.py

"""模型列表与模型解析的验证。

核心诉求：不再使用 gemini-webapi 已废弃的硬编码 ``Model`` 枚举，而是读账号初始化时
真实探测到的模型（``client.list_models()`` / ``client.resolve_model()``）。

为了真实地验证集成路径，这里给**真实的** ``GeminiClient`` 注入 ``_model_registry``，
再让它去做 list/resolve，不联网。
"""

import os
import sys
import tempfile
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-MODEL"
os.environ["SECURE_1PSIDTS"] = "TS-MODEL"
os.environ["GEMINI_COOKIE_PATH"] = tempfile.mkdtemp(prefix="gs-model-")
os.environ["SESSION_HEALTH_INTERVAL"] = "3600"

from gemini_webapi import GeminiClient as RealGeminiClient  # noqa: E402
from gemini_webapi.constants import Model  # noqa: E402
from gemini_webapi.types import AvailableModel  # 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 ""))


# 模拟账号真实探测到的模型：注意名字带版本号，且有一个不可选用
ACCOUNT_MODELS = [
	AvailableModel(
		model_id="9d8ca3786ebdfbea",
		model_name="gemini-3-pro",
		display_name="Gemini 3 Pro",
		description="最强推理模型",
		capacity=1,
		capacity_field=12,
		model_number=3,
		aliases=["3-pro", "gemini-3-pro-preview"],
	),
	AvailableModel(
		model_id="fbb127bbb056c959",
		model_name="gemini-3-flash",
		display_name="Gemini 3 Flash",
		description="速度与质量兼顾",
		capacity=1,
		capacity_field=12,
		model_number=1,
		aliases=["3-flash"],
	),
	AvailableModel(
		model_id="cf41b0e0dd7d53e5",
		model_name="gemini-3-flash-lite",
		display_name="Gemini 3 Flash Lite",
		description="最轻量",
		capacity=1,
		capacity_field=12,
		model_number=6,
		is_available=False,
	),
]


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):
		self.text = "OK"
		self.thoughts = None
		self.images = []
		self.metadata = ["cid-model"]


class FakeClient:
	"""模型相关能力直接委托给真实的 GeminiClient，其余用假实现避免联网。"""

	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.generate_calls = []

		self._real = RealGeminiClient()
		self._real._model_registry = {m.model_id: m for m in ACCOUNT_MODELS}

	def list_models(self):
		return self._real.list_models()

	def resolve_model(self, name):
		return self._real.resolve_model(name)

	async def init(self, **kwargs):
		self._running = True
		self.client = object()
		self.auto_refresh = kwargs.get("auto_refresh", False)
		if self.auto_refresh:
			import asyncio

			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):
		self.generate_calls.append(kwargs)
		return FakeResponse()

	async def generate_content_stream(self, prompt, **kwargs):
		self.generate_calls.append(kwargs)
		yield FakeResponse()

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

	async def delete_chat(self, cid):
		return None

	async def start_auto_refresh(self):
		import asyncio

		while self._running:
			await asyncio.sleep(3600)

	async def start_activity_watchdog(self):
		import asyncio

		while self._running:
			await asyncio.sleep(3600)

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


def main():
	gs.GeminiClient = FakeClient

	import main as app_main
	from fastapi.testclient import TestClient

	# ---------------------------------------------------------- 纯函数层
	probe = FakeClient()
	check("1a 解析 model_name 命中", app_main.resolve_model_for_client(probe, "gemini-3-pro").model_name == "gemini-3-pro")
	check("1b 解析 display_name 命中", app_main.resolve_model_for_client(probe, "Gemini 3 Flash").model_name == "gemini-3-flash")
	check("1c 解析别名命中", app_main.resolve_model_for_client(probe, "3-flash").model_name == "gemini-3-flash")
	check("1d 解析 model_id 命中", app_main.resolve_model_for_client(probe, "9d8ca3786ebdfbea").model_name == "gemini-3-pro")

	unknown = app_main.resolve_model_for_client(probe, "totally-unknown-model-xyz")
	check("2a 完全未知的模型返回 None（不再回退到废弃枚举）", unknown is None, repr(unknown))
	check("2b 未知模型不会返回 Model 枚举", not isinstance(unknown, Model))

	# gemini-webapi 内置「版本无关匹配」：旧名字会被映射到账号里最接近的当前模型，
	# 这比静默改用默认模型更好，属于库的正确行为，这里锁定住别退化。
	legacy = app_main.resolve_model_for_client(probe, "gemini-1.5-pro")
	check("2c 旧模型名按版本无关规则映射到当前 pro 模型",
	      isinstance(legacy, AvailableModel) and legacy.model_name == "gemini-3-pro", repr(legacy))
	check("2d 空模型名返回 None", app_main.resolve_model_for_client(probe, "") is None)
	check("2e client 为 None 时返回 None", app_main.resolve_model_for_client(None, "gemini-3-pro") is None)

	# 扩展思考标记的剥离：必须在 resolve 之前做，否则那个名字它不认识，会回退到默认模型
	check(
		"2f -thinking 后缀被剥离",
		app_main.split_extended_thinking_marker("gemini-3.8-flash-thinking") == ("gemini-3.8-flash", True),
	)
	check(
		"2g @extended 后缀被剥离",
		app_main.split_extended_thinking_marker("gemini-3.8-flash@extended") == ("gemini-3.8-flash", True),
	)
	check(
		"2h 无标记时原样返回",
		app_main.split_extended_thinking_marker("gemini-3.8-flash") == ("gemini-3.8-flash", False),
	)
	check("2i 大小写不敏感", app_main.split_extended_thinking_marker("X-THINKING") == ("X", True))
	check("2j 空值不报错", app_main.split_extended_thinking_marker("") == ("", False))
	check(
		"2k 剥离后的名字仍能正常解析",
		getattr(app_main.resolve_model_for_client(probe, "gemini-3-pro"), "model_name", None) == "gemini-3-pro",
	)

	with TestClient(app_main.app) as client:
		deadline = time.time() + 20
		while time.time() < deadline and not app_main.session_manager.status()["initialized"]:
			time.sleep(0.1)
		check("3 会话已初始化", app_main.session_manager.status()["initialized"] is True)

		# ---------------------------------------------------------- /v1/models
		res = client.get("/v1/models")
		check("4a /v1/models 返回 200", res.status_code == 200, str(res.status_code))
		body = res.json()
		ids = [m["id"] for m in body["data"]]
		print(f"       账号模型: {ids}  source={body.get('source')}")

		check("4b 来源标记为 account", body.get("source") == "account", str(body.get("source")))
		base_ids = [i for i in ids if not i.endswith("-thinking")]
		thinking_ids = [i for i in ids if i.endswith("-thinking")]
		check("4c 列出真实模型（带版本号）", base_ids == ["gemini-3-pro", "gemini-3-flash"], str(ids))
		check("4h 每个模型都有扩展思考变体", thinking_ids == ["gemini-3-pro-thinking", "gemini-3-flash-thinking"], str(ids))
		check("4i 扩展思考变体标明了 extends", body["data"][1].get("extends") == "gemini-3-pro", str(body["data"][1]))
		check("4d 过滤掉不可选用的模型", not any("flash-lite" in i for i in ids), str(ids))
		check("4e 不再出现废弃枚举的模型", not any("plus" in i or "advanced" in i for i in ids), str(ids))
		check("4f 带上了 display_name", body["data"][0].get("display_name") == "Gemini 3 Pro", str(body["data"][0]))
		check("4g 带上了 description", bool(body["data"][0].get("description")))

		# ---------------------------------------------------------- 客户端未就绪时的兜底
		saved = app_main.session_manager.client
		app_main.session_manager.client = None
		try:
			body_fb = client.get("/v1/models").json()
			check("5a 客户端未就绪时来源为 fallback", body_fb.get("source") == "fallback", str(body_fb.get("source")))
			check("5b 兜底列表非空", len(body_fb["data"]) > 0)
		finally:
			app_main.session_manager.client = saved

		# ---------------------------------------------------------- 请求里的模型被真实解析
		live = app_main.session_manager.client
		live.generate_calls.clear()
		res = client.post("/v1/chat/completions", json={
			"model": "gemini-3-pro",
			"messages": [{"role": "user", "content": "hi"}],
		})
		check("6a 指定模型时请求成功", res.status_code == 200, f"{res.status_code} {res.text[:160]}")
		passed_model = live.generate_calls[-1].get("model")
		check("6b 传给 generate_content 的是 AvailableModel", isinstance(passed_model, AvailableModel), repr(passed_model))
		check("6c 模型正是请求的那个", getattr(passed_model, "model_name", None) == "gemini-3-pro", repr(passed_model))
		check("6d 不是废弃的 Model 枚举", not isinstance(passed_model, Model))

		# ---------------------------------------------------------- 未知模型 -> 交给 Gemini 默认
		live.generate_calls.clear()
		res = client.post("/v1/chat/completions", json={
			"model": "totally-unknown-model-xyz",
			"messages": [{"role": "user", "content": "hi"}],
		})
		check("7a 未知模型时请求仍成功", res.status_code == 200, f"{res.status_code} {res.text[:160]}")
		check("7b 未知模型时不发送 model 参数（让 Gemini 用默认）", "model" not in live.generate_calls[-1],
		      str(live.generate_calls[-1].keys()))

		# ---------------------------------------------------------- 扩展思考（长输出）
		live.generate_calls.clear()
		res = client.post("/v1/chat/completions", json={
			"model": "gemini-3-pro-thinking",
			"messages": [{"role": "user", "content": "hi"}],
		})
		check("7c 带 -thinking 的请求成功", res.status_code == 200, f"{res.status_code} {res.text[:140]}")
		call = live.generate_calls[-1]
		check("7d 传了 extended_thinking=True", call.get("extended_thinking") is True, str(call.keys()))
		check(
			"7e -thinking 后缀已剥离，模型正确解析",
			getattr(call.get("model"), "model_name", None) == "gemini-3-pro",
			repr(call.get("model")),
		)

		live.generate_calls.clear()
		client.post("/v1/chat/completions", json={
			"model": "gemini-3-pro",
			"messages": [{"role": "user", "content": "hi"}],
		})
		check(
			"7f 普通模型不开启扩展思考",
			"extended_thinking" not in live.generate_calls[-1],
			str(live.generate_calls[-1].keys()),
		)

	# ---------------------------------------------------------- 对外 id 要带上版本号
	# 账号真实返回的形态：model_name 没有版本号，版本只在 display_name 里
	versioned = FakeClient()
	versioned._real._model_registry = {
		"aaa111": AvailableModel(
			model_id="aaa111", model_name="gemini-flash", display_name="3.8 Flash",
			description="", capacity=1, capacity_field=12, model_number=1,
			aliases=["gemini-3.8-flash", "flash"],
		),
		"bbb222": AvailableModel(
			model_id="bbb222", model_name="gemini-pro", display_name="3.1 Pro",
			description="", capacity=1, capacity_field=12, model_number=3,
			aliases=["gemini-3.1-pro", "pro"],
		),
	}
	flash = versioned._real._model_registry["aaa111"]
	pro = versioned._real._model_registry["bbb222"]

	id_flash = app_main.public_model_id(versioned, flash)
	id_pro = app_main.public_model_id(versioned, pro)
	print(f"       对外 id: {id_flash} / {id_pro}")

	check("8a 对外 id 带上版本号（flash）", id_flash == "gemini-3.8-flash", id_flash)
	check("8b 对外 id 带上版本号（pro）", id_pro == "gemini-3.1-pro", id_pro)
	check("8c 新 id 能精确解析回同一个模型", versioned.resolve_model(id_flash).model_id == "aaa111")
	check("8d 新 id 不会串到别的模型", versioned.resolve_model(id_pro).model_id == "bbb222")

	# 库的版本无关匹配（MODEL_PREFIX_RE = ^gemini-(?:\d+(?:\.\d+)?-)?）会把
	# gemini-<任意版本>-<档位> 归一化成 <档位>，所以带版本号的 id 天然能解析回来。
	orphan = AvailableModel(
		model_id="ccc333", model_name="gemini-lite", display_name="9.9 Lite",
		description="", capacity=1, capacity_field=12, model_number=6,
	)
	versioned._real._model_registry["ccc333"] = orphan
	check("8e 即使 model_name 无版本号，也能生成带版本号的 id",
	      app_main.public_model_id(versioned, orphan) == "gemini-9.9-lite",
	      app_main.public_model_id(versioned, orphan))
	check("8e2 且能解析回同一模型",
	      versioned.resolve_model("gemini-9.9-lite").model_id == "ccc333")

	# 真正需要兜底的场景：display_name 的 slug 撞上了**另一个**模型，
	# 此时必须退回 model_name，否则请求会被打到别的模型上。
	collide = AvailableModel(
		model_id="fff666", model_name="gemini-flash-lite-x", display_name="Flash",
		description="", capacity=1, capacity_field=12, model_number=6,
	)
	versioned._real._model_registry["fff666"] = collide
	check("8f slug 撞上别的模型时退回 model_name",
	      app_main.public_model_id(versioned, collide) == "gemini-flash-lite-x",
	      app_main.public_model_id(versioned, collide))
	check("8g display_name 为空时退回 model_name",
	      app_main.public_model_id(versioned, AvailableModel(
		      model_id="ddd444", model_name="gemini-x", display_name="",
		      description="", capacity=1, capacity_field=12, model_number=1)) == "gemini-x")
	check("8h 与 model_name 相同时不画蛇添足",
	      app_main.public_model_id(versioned, AvailableModel(
		      model_id="eee555", model_name="gemini-3.5-flash-lite", display_name="3.5 Flash-Lite",
		      description="", capacity=1, capacity_field=12, model_number=6)) == "gemini-3.5-flash-lite")

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