119 lines
4.9 KiB
Python
119 lines
4.9 KiB
Python
"""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 关闭）
|