sentinel-home-ai/fam-edge/tests/test_person_identifier.py

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