Files
sentinel-home-ai/fam-edge/tests/test_gemini_adapter.py
ericwyuan 708b6365f4 feat(fam-edge): Gemini key 分配改真随机 + 按项目名标注统计
用户明确要求: 每次调用应该是 1/n 概率的真随机抽取,而不是上一版的顺序轮转
(round-robin 在当前单消费者串行处理场景下已经是数学最优均匀分配,但用户想要
"数组里随机取一个,每次都是四分之一概率"这种更直观的随机语义)。_rotated_keys()
去掉持久的轮转下标,改用 random.randrange 每次调用随机选起点,仍保留"起点 key
全部模型失败就级联到下一个 key"的兜底逻辑不变。

同时新增 key_labels 配置:4 个 key 分别对应用户在 Google Cloud 上开的 4 个项目
(智能摄像头-1/2/3/4),model_calls 统计现在按项目名而不是泛化的 key1/key2
展示,ModelStats 页面的 Key 列/拆分逻辑同步改为按最后一个 "·" 分隔(原来写死
按 "·key" 前缀查找,项目名场景下不适用)。

新增/重写 7 个单测覆盖:随机起点级联顺序、randrange 调用范围正确性、大样本
随机分布均匀性、key_labels 默认值/自定义值/key 被丢弃时的对齐。
2026-08-22 07:37:33 +08:00

111 lines
4.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from fam_edge.model_adapters.gemini_adapter import GeminiAdapter
def _cfg(**overrides):
base = {
"provider": "gemini",
"model_name": "gemini-flash-latest",
"api_key": "key-primary",
"circuit_breaker": {"enabled": False},
}
base.update(overrides)
return base
def test_single_key_backward_compat():
a = GeminiAdapter(_cfg())
assert a.api_keys == ["key-primary"]
assert a.api_key == "key-primary"
def test_extra_keys_appended_in_order():
a = GeminiAdapter(_cfg(extra_api_keys=["key-2", "key-3", "key-4"]))
assert a.api_keys == ["key-primary", "key-2", "key-3", "key-4"]
def test_extra_keys_dedup_against_primary():
a = GeminiAdapter(_cfg(extra_api_keys=["key-primary", "key-2"]))
assert a.api_keys == ["key-primary", "key-2"]
def test_missing_env_var_keys_are_dropped():
"""extra_api_keys 里未设置的 ${ENV_VAR} 解析为空字符串,不应该混进 api_keys 列表。"""
a = GeminiAdapter(_cfg(extra_api_keys=["${SOME_UNSET_GEMINI_KEY_VAR}", "key-2"]))
assert a.api_keys == ["key-primary", "key-2"]
def test_no_keys_at_all():
a = GeminiAdapter(_cfg(api_key=""))
assert a.api_keys == []
assert a.api_key == ""
def test_key_labels_default_to_generic_when_not_configured():
a = GeminiAdapter(_cfg(extra_api_keys=["key-2", "key-3"]))
assert a.key_labels == ["key1", "key2", "key3"]
def test_key_labels_use_configured_project_names():
a = GeminiAdapter(_cfg(
extra_api_keys=["key-2", "key-3", "key-4"],
key_labels=["智能摄像头-1", "智能摄像头-2", "智能摄像头-3", "智能摄像头-4"],
))
assert a.key_labels == ["智能摄像头-1", "智能摄像头-2", "智能摄像头-3", "智能摄像头-4"]
def test_key_labels_stay_aligned_when_a_key_is_dropped():
"""如果某个 ${ENV_VAR} 没设置被丢弃,剩下的 key_labels 要跟着剩下的 key 对齐,
不能因为下标错位把别的项目名安到错的 key 上。"""
a = GeminiAdapter(_cfg(
extra_api_keys=["${SOME_UNSET_GEMINI_KEY_VAR}", "key-2"],
key_labels=["智能摄像头-1", "智能摄像头-2", "智能摄像头-3"],
))
assert a.api_keys == ["key-primary", "key-2"]
assert a.key_labels == ["智能摄像头-1", "智能摄像头-3"]
def test_rotated_keys_cascades_from_random_start(monkeypatch):
"""起点由 random.randrange 决定;固定住随机数就能验证级联顺序是"从起点绕一圈""""
a = GeminiAdapter(_cfg(extra_api_keys=["key-2", "key-3", "key-4"]))
monkeypatch.setattr(
"fam_edge.model_adapters.gemini_adapter.random.randrange", lambda n: 2)
order = a._rotated_keys()
assert [k for _, k in order] == ["key-3", "key-4", "key-primary", "key-2"]
assert [i for i, _ in order] == [2, 3, 0, 1]
def test_rotated_keys_uses_full_key_range(monkeypatch):
"""random.randrange 的调用范围必须是 len(api_keys),否则会漏掉某些 key 永远选不到。"""
a = GeminiAdapter(_cfg(extra_api_keys=["key-2", "key-3", "key-4"]))
seen_n = []
def fake_randrange(n):
seen_n.append(n)
return 0
monkeypatch.setattr(
"fam_edge.model_adapters.gemini_adapter.random.randrange", fake_randrange)
a._rotated_keys()
assert seen_n == [4]
def test_rotated_keys_starts_roughly_uniform_over_many_calls():
"""核心诉求: 每次调用应该是真随机(每个 key 命中概率均等 1/n而不是像旧的
round-robin 那样顺序推进——用大样本统计每个 key 被选为起点的频率,应该接近
1/4且不应该出现某个 key 明显被冷落或独占(对应此前"全部流量压在同一个
key"的 bug"""
a = GeminiAdapter(_cfg(extra_api_keys=["key-2", "key-3", "key-4"]))
n_trials = 4000
counts = {"key-primary": 0, "key-2": 0, "key-3": 0, "key-4": 0}
for _ in range(n_trials):
counts[a._rotated_keys()[0][1]] += 1
for key, c in counts.items():
share = c / n_trials
assert 0.20 <= share <= 0.30, f"{key} 起点占比 {share} 明显偏离 1/4"
def test_rotated_keys_single_key_never_errors():
a = GeminiAdapter(_cfg())
for _ in range(3):
assert a._rotated_keys() == [(0, "key-primary")]