yawan-workbench/server/tests/test_llm_failover.py

"""llm.chat failover 接线端到端测试（mock OpenAI 客户端，不依赖网络）。"""
from types import SimpleNamespace
from unittest.mock import patch, MagicMock

import app.llm as llm_mod
from app.llm import chat


def _provider_row(pid, name, models):
    return SimpleNamespace(id=pid, name=name, models=list(models),
                           default_model=models[0] if models else "", enabled=True,
                           base_url="https://mock.local/v1", api_key="mock-key")


class FakeSession:
    """最小 session mock：get 返回 None（无 campaign/run 归属），exec 供 failover 查询。"""

    def __init__(self, providers):
        self._providers = providers

    def get(self, model, row_id):
        return None

    def exec(self, *_a, **_k):
        rows = list(self._providers.values())
        return SimpleNamespace(all=lambda: rows,
                               first=lambda: None)

    def commit(self):
        pass

    def add(self, _obj):
        pass


def _fake_resp(text="你好，正文内容。"):
    resp = MagicMock()
    resp.choices = [SimpleNamespace(message=SimpleNamespace(
        content=text, tool_calls=None))]
    resp.usage = SimpleNamespace(prompt_tokens=10, completion_tokens=20)
    return resp


def test_chat_failover_on_auth_error():
    """主 Provider 认证失败 → 自动切备用 Provider 成功。"""
    providers = {1: _provider_row(1, "主Provider", ["deepseek-chat"]),
                 2: _provider_row(2, "备Provider", ["deepseek-chat"])}
    calls = {"n": 0}

    def fake_create(**request):
        calls["n"] += 1
        if calls["n"] == 1:
            raise RuntimeError("Error code: 401 - {'error': {'message': 'Authentication Fails'}}")
        return _fake_resp()

    session = FakeSession(providers)
    with patch.object(llm_mod, "get_provider", side_effect=lambda s, pid: providers[pid]), \
         patch.object(llm_mod, "require_ai_points", return_value=None), \
         patch.object(llm_mod, "_log", return_value=True), \
         patch.object(llm_mod, "record_usage_event", return_value=SimpleNamespace(id=1)), \
         patch.object(llm_mod.OpenAI, "chat") as mock_chat:
        mock_chat.completions.create.side_effect = fake_create
        result = chat(session, [{"role": "user", "content": "测试"}], 1, "deepseek-chat",
                      user_id=1, organization_id=1)
    assert result["status"] == "ok"
    assert result["text"] == "你好，正文内容。"
    assert result["failover"]["from_provider"] == "主Provider"
    assert result["failover"]["from_error"] == "auth"
    assert calls["n"] == 2


def test_chat_no_failover_on_other_error():
    """不可切换错误（other）不触发 failover。"""
    providers = {1: _provider_row(1, "主Provider", ["deepseek-chat"]),
                 2: _provider_row(2, "备Provider", ["deepseek-chat"])}
    calls = {"n": 0}

    def fake_create(**request):
        calls["n"] += 1
        raise RuntimeError("some mysterious failure")

    session = FakeSession(providers)
    with patch.object(llm_mod, "get_provider", side_effect=lambda s, pid: providers[pid]), \
         patch.object(llm_mod, "require_ai_points", return_value=None), \
         patch.object(llm_mod, "_log", return_value=True), \
         patch.object(llm_mod, "record_usage_event", return_value=SimpleNamespace(id=1)), \
         patch.object(llm_mod.OpenAI, "chat") as mock_chat:
        mock_chat.completions.create.side_effect = fake_create
        result = chat(session, [{"role": "user", "content": "测试"}], 1, "deepseek-chat",
                      user_id=1, organization_id=1, max_retries=0)
    assert result["status"] == "failed"
    assert "failover" not in result
    assert calls["n"] == 1


def test_chat_failover_all_providers_down():
    """备用 Provider 也失败 → 返回原始失败并注明主 Provider 错误。"""
    providers = {1: _provider_row(1, "主Provider", ["deepseek-chat"]),
                 2: _provider_row(2, "备Provider", ["deepseek-chat"])}
    calls = {"n": 0}

    def fake_create(**request):
        calls["n"] += 1
        if calls["n"] == 1:
            raise RuntimeError("Error code: 429 - rate limit exceeded")
        raise RuntimeError("Error code: 401 - Authentication Fails")

    session = FakeSession(providers)
    with patch.object(llm_mod, "get_provider", side_effect=lambda s, pid: providers[pid]), \
         patch.object(llm_mod, "require_ai_points", return_value=None), \
         patch.object(llm_mod, "_log", return_value=True), \
         patch.object(llm_mod, "record_usage_event", return_value=SimpleNamespace(id=1)), \
         patch.object(llm_mod.OpenAI, "chat") as mock_chat:
        mock_chat.completions.create.side_effect = fake_create
        result = chat(session, [{"role": "user", "content": "测试"}], 1, "deepseek-chat",
                      user_id=1, organization_id=1, max_retries=0)
    assert result["status"] == "failed"
    assert calls["n"] == 2  # 主 1 次 + 备 1 次（备的重试被 max_retries=0 关闭）