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

import json

import pytest

from fam_edge.qa import QAOrchestrator


def _orchestrator(monkeypatch, cfg=None, token='tok'):
    ai_gateway_cfg = {"base_url": "http://127.0.0.1:5100", "token": token, "timeout": 5}
    if cfg:
        ai_gateway_cfg.update(cfg)
    monkeypatch.setattr(
        "fam_edge.qa.load_config", lambda: {"ai_gateway": ai_gateway_cfg})
    return QAOrchestrator()


def test_init_reads_base_url_and_token_from_config(monkeypatch):
    qa = _orchestrator(monkeypatch, {"base_url": "http://example:5100/"}, token='secret')
    assert qa.base_url == "http://example:5100"
    assert qa.token == 'secret'


def test_init_resolves_token_from_env_var(monkeypatch):
    monkeypatch.setenv("MY_GATEWAY_TOKEN", "resolved-secret")
    qa = _orchestrator(monkeypatch, token='${MY_GATEWAY_TOKEN}')
    assert qa.token == 'resolved-secret'


def test_init_defaults_base_url_when_unconfigured(monkeypatch):
    monkeypatch.setattr("fam_edge.qa.load_config", lambda: {})
    qa = QAOrchestrator()
    assert qa.base_url == "http://127.0.0.1:5100"


class _FakeResp:
    """模拟 requests.Response：非流式用 status_code/json()/text，
    流式额外提供 iter_lines()（逐行 yield，跟真实 SSE 消费方式一致）。"""

    def __init__(self, status_code=200, payload=None, text='', lines=None):
        self.status_code = status_code
        self._payload = payload
        self.text = text
        self._lines = lines if lines is not None else []
        self.encoding = None

    def json(self):
        return self._payload

    def iter_lines(self, decode_unicode=True):
        for line in self._lines:
            yield line


def _capture_post(monkeypatch, resp):
    calls = []

    def fake_post(url, headers=None, json=None, timeout=None, stream=False):
        calls.append({"url": url, "headers": headers, "json": json,
                      "timeout": timeout, "stream": stream})
        return resp

    monkeypatch.setattr("fam_edge.qa.requests.post", fake_post)
    return calls


def test_run_qa_success(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(payload={"choices": [{"message": {"content": "你好"}}],
                              "provider": "nvidia"})
    calls = _capture_post(monkeypatch, resp)
    answer, provider = qa.run_qa("hi", max_tokens=100)
    assert answer == "你好"
    assert provider == "nvidia"
    assert calls[0]["json"] == {"messages": [{"role": "user", "content": "hi"}],
                                "max_tokens": 100, "stream": False}
    assert calls[0]["headers"]["Authorization"] == "Bearer tok"


def test_run_qa_non_200_returns_none(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(status_code=503, text='{"error":{"message":"所有模型均不可用"}}')
    _capture_post(monkeypatch, resp)
    answer, provider = qa.run_qa("hi")
    assert answer is None
    assert provider is None


def test_run_qa_empty_answer_returns_none(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(payload={"choices": [{"message": {"content": ""}}], "provider": "gemini"})
    _capture_post(monkeypatch, resp)
    answer, provider = qa.run_qa("hi")
    assert answer is None
    assert provider is None


def test_run_qa_connection_error_returns_none(monkeypatch):
    qa = _orchestrator(monkeypatch)

    def _raise(*args, **kwargs):
        raise ConnectionError("boom")

    monkeypatch.setattr("fam_edge.qa.requests.post", _raise)
    answer, provider = qa.run_qa("hi")
    assert answer is None
    assert provider is None


def _sse_lines(events):
    lines = []
    for e in events:
        lines.append(f"data: {json.dumps(e, ensure_ascii=False)}")
    lines.append("data: [DONE]")
    return lines


def test_run_qa_stream_single_provider_success(monkeypatch):
    qa = _orchestrator(monkeypatch)
    lines = _sse_lines([
        {"provider": "nvidia", "choices": [{"delta": {"content": "你"}}]},
        {"provider": "nvidia", "choices": [{"delta": {"content": "好"}}]},
        {"provider": "nvidia", "choices": [{"delta": {}}]},
    ])
    resp = _FakeResp(lines=lines)
    _capture_post(monkeypatch, resp)
    events = list(qa.run_qa_stream("hi"))
    types = [e["type"] for e in events]
    assert types == ["provider_trying", "chunk", "chunk", "done"]
    assert events[1]["text"] == "你"
    assert events[2]["text"] == "好"
    assert events[-1]["provider"] == "nvidia"


def test_run_qa_stream_emits_provider_trying_once_per_change(monkeypatch):
    """provider 字段没变化时不该重复吐 provider_trying。"""
    qa = _orchestrator(monkeypatch)
    lines = _sse_lines([
        {"provider": "nvidia", "choices": [{"delta": {"content": "a"}}]},
        {"provider": "nvidia", "choices": [{"delta": {"content": "b"}}]},
    ])
    resp = _FakeResp(lines=lines)
    _capture_post(monkeypatch, resp)
    events = list(qa.run_qa_stream("hi"))
    trying = [e for e in events if e["type"] == "provider_trying"]
    assert len(trying) == 1
    assert trying[0]["provider"] == "nvidia"


def test_run_qa_stream_no_chunks_yields_all_failed(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(lines=["data: [DONE]"])
    _capture_post(monkeypatch, resp)
    events = list(qa.run_qa_stream("hi"))
    assert events == [{"type": "all_failed"}]


def test_run_qa_stream_non_200_yields_all_failed(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(status_code=503, text='{"error":{"message":"所有模型均不可用"}}')
    _capture_post(monkeypatch, resp)
    events = list(qa.run_qa_stream("hi"))
    assert events == [{"type": "all_failed"}]


def test_run_qa_stream_connection_error_yields_all_failed(monkeypatch):
    qa = _orchestrator(monkeypatch)

    def _raise(*args, **kwargs):
        raise ConnectionError("boom")

    monkeypatch.setattr("fam_edge.qa.requests.post", _raise)
    events = list(qa.run_qa_stream("hi"))
    assert events == [{"type": "all_failed"}]


def test_run_qa_stream_error_chunk_stops_and_uses_partial_output(monkeypatch):
    """已经吐出过内容后遇到错误块：按"至少吐出过一块就算 done"处理，不是 all_failed。"""
    qa = _orchestrator(monkeypatch)
    lines = [
        f"data: {json.dumps({'provider': 'gemini', 'choices': [{'delta': {'content': '先吐'}}]}, ensure_ascii=False)}",
        f"data: {json.dumps({'error': {'message': 'boom'}}, ensure_ascii=False)}",
    ]
    resp = _FakeResp(lines=lines)
    _capture_post(monkeypatch, resp)
    events = list(qa.run_qa_stream("hi"))
    types = [e["type"] for e in events]
    assert types == ["provider_trying", "chunk", "done"]
    assert events[-1]["provider"] == "gemini"


def test_run_qa_stream_sets_stream_true_and_utf8_encoding(monkeypatch):
    qa = _orchestrator(monkeypatch)
    resp = _FakeResp(lines=["data: [DONE]"])
    calls = _capture_post(monkeypatch, resp)
    list(qa.run_qa_stream("hi", max_tokens=222))
    assert calls[0]["json"]["stream"] is True
    assert calls[0]["json"]["max_tokens"] == 222
    assert calls[0]["stream"] is True
    assert resp.encoding == 'utf-8'