import os import pytest from fam_edge.person_identifier import PersonIdentifier def _cfg(ref_dir, **overrides): base = { "enabled": True, "ref_dir": ref_dir, "max_ref_per_person": 6, "min_call_interval_sec": 0, # 测试不需要真实限速,避免拖慢用例 "nvidia": {"api_key": "nvkey", "model_name": "nvidia/test", "timeout": 30, "max_retries": 2, "retry_backoff_sec": 0.01}, "gemini": {"api_key": "gkey1", "model_name": "gemini-flash-lite-latest", "timeout": 30, "max_retries": 2, "retry_backoff_sec": 0.01}, } base.update(overrides) return base def _write_refs(tmp_path, grandpa=2, dad=2): for person, n in (("爷爷", grandpa), ("爸爸", dad)): d = tmp_path / person d.mkdir(parents=True, exist_ok=True) for i in range(n): (d / f"{i:02d}.jpg").write_bytes(b"fakejpegbytes") class _FakeResp: def __init__(self, status_code=200, payload=None): self.status_code = status_code self._payload = payload or {} def json(self): return self._payload @pytest.fixture(autouse=True) def _no_real_sleep(monkeypatch): """全部用例都不需要真的睡(限速/退避都测计数和结果,不测真实耗时)。""" monkeypatch.setattr("fam_edge.person_identifier.time.sleep", lambda s: None) def test_has_references_false_when_dirs_missing(tmp_path): pi = PersonIdentifier(_cfg(str(tmp_path / "refs"))) assert pi.has_references() is False def test_has_references_true_when_both_present(tmp_path): _write_refs(tmp_path, grandpa=3, dad=2) pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.has_references() is True def test_has_references_false_when_only_one_person_has_refs(tmp_path): (tmp_path / "爷爷").mkdir(parents=True) (tmp_path / "爷爷" / "01.jpg").write_bytes(b"x") pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.has_references() is False def test_extract_json_person_clean_json(): pi = PersonIdentifier(_cfg("/nonexistent")) assert pi._extract_json_person('{"person":"爷爷"}') == '爷爷' assert pi._extract_json_person('{"person":"爸爸"}') == '爸爸' def test_extract_json_person_no_match_returns_none(): pi = PersonIdentifier(_cfg("/nonexistent")) assert pi._extract_json_person('不知道是谁') is None assert pi._extract_json_person('') is None def test_extract_json_person_both_mentioned_uses_json_field(): """核心诉求: 模型有时会把参考图说明也复述一遍,回复里两个名字都出现—— 这时候不能瞎猜,要从 JSON 的 person 字段里精确取,取不到就返回 None。""" pi = PersonIdentifier(_cfg("/nonexistent")) text = '参考图1是爷爷,参考图5是爸爸。{"person":"爸爸"}' assert pi._extract_json_person(text) == '爸爸' def test_extract_json_person_both_mentioned_no_json_field_returns_none(): pi = PersonIdentifier(_cfg("/nonexistent")) text = '这个人可能是爷爷,也可能是爸爸,不太确定' assert pi._extract_json_person(text) is None def test_classify_returns_none_when_disabled(tmp_path): _write_refs(tmp_path) pi = PersonIdentifier(_cfg(str(tmp_path), enabled=False)) assert pi.classify_adult_male(b"crop") is None def test_classify_returns_none_when_no_references(tmp_path): pi = PersonIdentifier(_cfg(str(tmp_path / "empty"))) assert pi.classify_adult_male(b"crop") is None def test_classify_falls_back_to_gemini_when_nvidia_unavailable(tmp_path, monkeypatch): """openai SDK 未安装时 NVIDIA 路径应该静默跳过(不报错),落到 Gemini。""" _write_refs(tmp_path) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", None) def fake_post(url, json=None, timeout=None): return _FakeResp(200, { "candidates": [{"content": {"parts": [{"text": '{"person":"爸爸"}'}]}}] }) monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.classify_adult_male(b"crop") == '爸爸' def _fake_openai_factory(reply_text=None, exc=None, fail_times=0): """构造一个假 OpenAI 客户端:先失败 fail_times 次再成功,或者一直抛 exc。""" state = {"calls": 0} class FakeMessage: content = reply_text class FakeChoice: message = FakeMessage() class FakeChatResp: choices = [FakeChoice()] class FakeCompletions: def create(self, **kwargs): state["calls"] += 1 if state["calls"] <= fail_times: raise (exc or RuntimeError("boom")) if exc and fail_times == 0: raise exc return FakeChatResp() class FakeChat: completions = FakeCompletions() class FakeOpenAI: def __init__(self, base_url=None, api_key=None): pass chat = FakeChat() return FakeOpenAI, state def test_classify_nvidia_success_skips_gemini(tmp_path, monkeypatch): _write_refs(tmp_path) FakeOpenAI, state = _fake_openai_factory(reply_text='{"person":"爷爷"}') monkeypatch.setattr("fam_edge.person_identifier.OpenAI", FakeOpenAI) gemini_called = {"n": 0} def fake_post(url, json=None, timeout=None): gemini_called["n"] += 1 return _FakeResp(200, {}) monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.classify_adult_male(b"crop") == '爷爷' assert gemini_called["n"] == 0 assert state["calls"] == 1 class _FakeHTTPError(Exception): def __init__(self, status_code): self.response = type("R", (), {"status_code": status_code})() def test_nvidia_retries_transient_error_then_succeeds(tmp_path, monkeypatch): """核心诉求: 429/503 这类瞬时故障要退避重试,不是第一次失败就放弃换 provider。""" _write_refs(tmp_path) FakeOpenAI, state = _fake_openai_factory( reply_text='{"person":"爸爸"}', exc=_FakeHTTPError(503), fail_times=1) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", FakeOpenAI) pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.classify_adult_male(b"crop") == '爸爸' assert state["calls"] == 2 # 第一次 503 失败重试一次后成功 def test_nvidia_gives_up_after_max_retries_falls_back_to_gemini(tmp_path, monkeypatch): _write_refs(tmp_path) FakeOpenAI, state = _fake_openai_factory(exc=_FakeHTTPError(503), fail_times=99) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", FakeOpenAI) def fake_post(url, json=None, timeout=None): return _FakeResp(200, { "candidates": [{"content": {"parts": [{"text": '{"person":"汤圆"}'}]}}] }) # 用一个不属于爷爷/爸爸的返回值只是为了确认真的调用到了 gemini 分支 monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) cfg = _cfg(str(tmp_path)) pi = PersonIdentifier(cfg) pi.classify_adult_male(b"crop") assert state["calls"] == pi.nvidia_max_retries # 重试到上限就放弃,不会无限重试 def test_nvidia_non_retryable_error_gives_up_immediately(tmp_path, monkeypatch): """核心诉求: 400 参数错误这类非瞬时故障,重试没有意义,应该立刻换下一个模型/provider, 不要浪费时间重试一个注定失败的请求。""" _write_refs(tmp_path) FakeOpenAI, state = _fake_openai_factory(exc=_FakeHTTPError(400), fail_times=99) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", FakeOpenAI) monkeypatch.setattr("fam_edge.person_identifier.requests.post", lambda *a, **k: _FakeResp(500, {"error": "down"})) pi = PersonIdentifier(_cfg(str(tmp_path))) pi.classify_adult_male(b"crop") assert state["calls"] == 1 # 400 不重试,一次就放弃这个模型 def test_nvidia_falls_through_model_chain(tmp_path, monkeypatch): """核心诉求: 第一个模型重试耗尽后,应该换模型链里的下一个型号再试,而不是 直接放弃整个 NVIDIA provider。""" _write_refs(tmp_path) calls = [] class FakeMessage: def __init__(self, content): self.content = content class FakeChoice: def __init__(self, content): self.message = FakeMessage(content) class FakeChatResp: def __init__(self, content): self.choices = [FakeChoice(content)] class FakeCompletions: def create(self, model, **kwargs): calls.append(model) if model == 'nvidia/model-a': raise _FakeHTTPError(503) return FakeChatResp('{"person":"爷爷"}') class FakeChat: completions = FakeCompletions() class FakeOpenAI: def __init__(self, base_url=None, api_key=None): pass chat = FakeChat() monkeypatch.setattr("fam_edge.person_identifier.OpenAI", FakeOpenAI) cfg = _cfg(str(tmp_path), nvidia={ "api_key": "nvkey", "model_name": "nvidia/model-a", "fallback_models": ["nvidia/model-b"], "timeout": 30, "max_retries": 2, "retry_backoff_sec": 0.01, }) pi = PersonIdentifier(cfg) assert pi.classify_adult_male(b"crop") == '爷爷' assert calls == ['nvidia/model-a', 'nvidia/model-a', 'nvidia/model-b'] def test_gemini_retries_transient_error_on_same_key(tmp_path, monkeypatch): _write_refs(tmp_path) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", None) calls = [] def fake_post(url, json=None, timeout=None): calls.append(url.split('key=')[-1]) if len(calls) == 1: return _FakeResp(503, {"error": {"code": 503}}) return _FakeResp(200, { "candidates": [{"content": {"parts": [{"text": '{"person":"媳妇"}'}]}}] }) monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) pi = PersonIdentifier(_cfg(str(tmp_path))) pi.classify_adult_male(b"crop") assert calls == ['gkey1', 'gkey1'] # 同一个 key 重试,不是立刻跳到下一个 key def test_classify_gemini_rotates_across_keys_after_retries_exhausted(tmp_path, monkeypatch): _write_refs(tmp_path) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", None) calls = [] def fake_post(url, json=None, timeout=None): key = url.split('key=')[-1] calls.append(key) if key == 'gkey1': return _FakeResp(429, {"error": {"code": 429}}) return _FakeResp(200, { "candidates": [{"content": {"parts": [{"text": '{"person":"爷爷"}'}]}}] }) monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) cfg = _cfg(str(tmp_path), gemini={ "api_key": "gkey1", "extra_api_keys": ["gkey2"], "model_name": "gemini-flash-lite-latest", "timeout": 30, "max_retries": 2, "retry_backoff_sec": 0.01, }) pi = PersonIdentifier(cfg) assert pi.classify_adult_male(b"crop") == '爷爷' assert calls == ['gkey1', 'gkey1', 'gkey2'] # gkey1 重试用尽才换 gkey2 def test_classify_both_providers_fail_returns_none(tmp_path, monkeypatch): """核心诉求: NVIDIA 和 Gemini 都失败时绝不能瞎猜,必须返回 None。""" _write_refs(tmp_path) monkeypatch.setattr("fam_edge.person_identifier.OpenAI", None) def fake_post(url, json=None, timeout=None): return _FakeResp(500, {"error": "boom"}) monkeypatch.setattr("fam_edge.person_identifier.requests.post", fake_post) pi = PersonIdentifier(_cfg(str(tmp_path))) assert pi.classify_adult_male(b"crop") is None def test_env_var_credentials_resolved(tmp_path): os.environ["TEST_NVIDIA_KEY_XYZ"] = "realkey" try: cfg = _cfg(str(tmp_path), nvidia={"api_key": "${TEST_NVIDIA_KEY_XYZ}"}) pi = PersonIdentifier(cfg) assert pi.nvidia_api_key == "realkey" finally: del os.environ["TEST_NVIDIA_KEY_XYZ"] def test_max_ref_per_person_limits_loaded_refs(tmp_path): _write_refs(tmp_path, grandpa=10, dad=10) pi = PersonIdentifier(_cfg(str(tmp_path), max_ref_per_person=3)) refs = pi._load_refs() assert len(refs['爷爷']) == 3 assert len(refs['爸爸']) == 3 def test_pace_sleeps_when_called_too_soon(tmp_path, monkeypatch): """核心诉求: 批量回填会短时间内密集调用,min_call_interval_sec 要真的限速, 不能形同虚设。""" _write_refs(tmp_path) slept = [] monkeypatch.setattr("fam_edge.person_identifier.time.sleep", lambda s: slept.append(s)) pi = PersonIdentifier(_cfg(str(tmp_path), min_call_interval_sec=5)) pi._last_call_at = __import__("time").time() # 刚刚调用过 pi._pace() assert slept and slept[0] > 0 def test_pace_no_sleep_when_interval_already_elapsed(tmp_path, monkeypatch): _write_refs(tmp_path) slept = [] monkeypatch.setattr("fam_edge.person_identifier.time.sleep", lambda s: slept.append(s)) pi = PersonIdentifier(_cfg(str(tmp_path), min_call_interval_sec=5)) pi._last_call_at = 0 # 很久以前 pi._pace() assert slept == []