209 lines
9.3 KiB
Python
209 lines
9.3 KiB
Python
"""模型 Key 兜底：402 识别、主 / 兜底选择、耗尽标记的持久化与自动恢复、Agent 用的包装模型。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime, timedelta
|
||
|
||
import pytest
|
||
|
||
from backend.browser import llm_factory, llm_keys
|
||
|
||
|
||
@pytest.fixture
|
||
def clean_keys(monkeypatch):
|
||
"""把 settings 表换成内存字典，并复位模块内缓存，测试之间互不影响。"""
|
||
store: dict[str, str] = {}
|
||
monkeypatch.setattr(llm_keys.db, "get_setting", lambda key, default="": store.get(key, default))
|
||
monkeypatch.setattr(llm_keys.db, "set_setting", lambda key, value: store.__setitem__(key, value))
|
||
monkeypatch.setattr(llm_keys, "_exhausted_at", None)
|
||
monkeypatch.setattr(llm_keys, "_loaded", True)
|
||
from backend.config import settings
|
||
|
||
monkeypatch.setattr(settings, "deepseek_api_key_fallback", "FALLBACK")
|
||
monkeypatch.setattr(settings, "dashscope_api_key_fallback", "")
|
||
monkeypatch.setattr(settings, "llm_primary_retry_hours", 1.0)
|
||
return store
|
||
|
||
|
||
class TestBalanceErrorDetection:
|
||
@pytest.mark.parametrize("error", [
|
||
"Error code: 402 - {'error': {'message': 'Insufficient Balance'}}",
|
||
"Insufficient Balance",
|
||
"Arrearage",
|
||
])
|
||
def test_recognised_from_the_message(self, error):
|
||
assert llm_keys.is_balance_error(error) is True
|
||
|
||
def test_recognised_from_the_status_code(self):
|
||
error = RuntimeError("something")
|
||
error.status_code = 402
|
||
assert llm_keys.is_balance_error(error) is True
|
||
|
||
@pytest.mark.parametrize("error", [
|
||
"Error code: 401 - Unauthorized",
|
||
"Error code: 429 - Rate limit",
|
||
"timeout waiting for selector",
|
||
"",
|
||
])
|
||
def test_other_failures_are_not_balance_errors(self, error):
|
||
# 限流、鉴权失败、页面超时都不该触发切 Key
|
||
assert llm_keys.is_balance_error(error) is False
|
||
|
||
|
||
class TestKeySelection:
|
||
def test_primary_is_used_while_it_works(self, clean_keys):
|
||
assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False)
|
||
|
||
def test_fallback_is_used_after_the_primary_is_marked_exhausted(self, clean_keys):
|
||
llm_keys.mark_primary_exhausted("402")
|
||
assert llm_keys.choose("deepseek", "PRIMARY") == ("FALLBACK", True)
|
||
|
||
def test_no_fallback_configured_keeps_using_the_primary(self, clean_keys, monkeypatch):
|
||
from backend.config import settings
|
||
|
||
monkeypatch.setattr(settings, "deepseek_api_key_fallback", "")
|
||
llm_keys.mark_primary_exhausted("402")
|
||
assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False)
|
||
|
||
def test_fallback_identical_to_the_primary_is_not_a_fallback(self, clean_keys, monkeypatch):
|
||
from backend.config import settings
|
||
|
||
monkeypatch.setattr(settings, "deepseek_api_key_fallback", "PRIMARY")
|
||
llm_keys.mark_primary_exhausted("402")
|
||
assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False)
|
||
|
||
def test_only_a_fallback_configured_is_still_usable(self, clean_keys):
|
||
key, using = llm_keys.choose("deepseek", "")
|
||
assert key == "FALLBACK" and using is True
|
||
|
||
def test_qwen_reads_its_own_fallback(self, clean_keys, monkeypatch):
|
||
from backend.config import settings
|
||
|
||
monkeypatch.setattr(settings, "dashscope_api_key_fallback", "QWEN_FALLBACK")
|
||
llm_keys.mark_primary_exhausted("402")
|
||
assert llm_keys.choose("qwen", "PRIMARY") == ("QWEN_FALLBACK", True)
|
||
|
||
|
||
class TestExhaustionMarker:
|
||
def test_marker_is_persisted(self, clean_keys):
|
||
llm_keys.mark_primary_exhausted("402")
|
||
assert clean_keys[llm_keys.EXHAUSTED_KEY]
|
||
assert llm_keys.primary_exhausted() is True
|
||
|
||
def test_clearing_restores_the_primary(self, clean_keys):
|
||
llm_keys.mark_primary_exhausted("402")
|
||
llm_keys.clear_primary_exhausted()
|
||
assert llm_keys.primary_exhausted() is False
|
||
assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False)
|
||
|
||
def test_primary_is_retried_after_the_retry_window(self, clean_keys, monkeypatch):
|
||
# 充值之后不用重启：过了重试周期下一次调用先试主 Key
|
||
monkeypatch.setattr(llm_keys, "_exhausted_at", datetime.now() - timedelta(hours=2))
|
||
assert llm_keys.primary_exhausted() is False
|
||
|
||
def test_marker_still_holds_inside_the_retry_window(self, clean_keys, monkeypatch):
|
||
monkeypatch.setattr(llm_keys, "_exhausted_at", datetime.now() - timedelta(minutes=10))
|
||
assert llm_keys.primary_exhausted() is True
|
||
|
||
def test_status_reports_what_the_ui_shows(self, clean_keys):
|
||
assert llm_keys.status("deepseek", "PRIMARY") == {
|
||
"fallback_configured": True, "using_fallback": False,
|
||
"primary_exhausted_at": None, "retry_hours": 1.0,
|
||
}
|
||
llm_keys.mark_primary_exhausted("402")
|
||
status = llm_keys.status("deepseek", "PRIMARY")
|
||
assert status["using_fallback"] is True and status["primary_exhausted_at"]
|
||
|
||
|
||
class FakeModel:
|
||
def __init__(self, name: str, fail: bool = False):
|
||
self.name = name
|
||
self.provider = "deepseek"
|
||
self.calls = 0
|
||
self.fail = fail
|
||
|
||
async def ainvoke(self, messages, output_format=None, **kwargs):
|
||
self.calls += 1
|
||
if self.fail:
|
||
from browser_use.llm.exceptions import ModelProviderError
|
||
|
||
raise ModelProviderError("Error code: 402 - {'error': {'message': 'Insufficient Balance'}}", model=self.name)
|
||
return f"answer-from-{self.name}"
|
||
|
||
|
||
class TestFallbackChatModel:
|
||
"""Agent 拿到的是这个包装对象：跑到一半余额耗尽也要接着跑完，不能让任务失败。"""
|
||
|
||
def _wrap(self, primary_fails: bool):
|
||
primary, fallback = FakeModel("primary", fail=primary_fails), FakeModel("fallback")
|
||
return llm_factory.FallbackChatModel(primary=primary, fallback=fallback, model="deepseek-chat"), primary, fallback
|
||
|
||
async def test_healthy_primary_is_used_and_the_fallback_stays_idle(self, clean_keys):
|
||
wrapper, primary, fallback = self._wrap(primary_fails=False)
|
||
assert await wrapper.ainvoke([], None) == "answer-from-primary"
|
||
assert primary.calls == 1 and fallback.calls == 0
|
||
|
||
async def test_402_switches_to_the_fallback_and_retries_the_same_call(self, clean_keys):
|
||
wrapper, primary, fallback = self._wrap(primary_fails=True)
|
||
assert await wrapper.ainvoke([], None) == "answer-from-fallback"
|
||
assert primary.calls == 1 and fallback.calls == 1
|
||
assert wrapper.using_fallback is True
|
||
assert llm_keys.primary_exhausted() is True
|
||
|
||
async def test_later_calls_go_straight_to_the_fallback(self, clean_keys):
|
||
wrapper, primary, _ = self._wrap(primary_fails=True)
|
||
await wrapper.ainvoke([], None)
|
||
await wrapper.ainvoke([], None)
|
||
assert primary.calls == 1 # 不会每一步都再撞一次主 Key
|
||
|
||
async def test_a_new_wrapper_starts_on_the_fallback_when_the_marker_is_set(self, clean_keys):
|
||
llm_keys.mark_primary_exhausted("402")
|
||
wrapper, primary, _ = self._wrap(primary_fails=True)
|
||
assert wrapper.using_fallback is True
|
||
assert await wrapper.ainvoke([], None) == "answer-from-fallback"
|
||
assert primary.calls == 0
|
||
|
||
async def test_a_successful_primary_call_clears_a_stale_marker(self, clean_keys):
|
||
llm_keys.mark_primary_exhausted("402")
|
||
llm_keys.clear_primary_exhausted() # 重试周期到了
|
||
wrapper, _, fallback = self._wrap(primary_fails=False)
|
||
await wrapper.ainvoke([], None)
|
||
assert llm_keys.primary_exhausted() is False and fallback.calls == 0
|
||
|
||
async def test_non_balance_errors_are_not_swallowed(self, clean_keys):
|
||
from browser_use.llm.exceptions import ModelProviderError
|
||
|
||
class Boom(FakeModel):
|
||
async def ainvoke(self, messages, output_format=None, **kwargs):
|
||
self.calls += 1
|
||
raise ModelProviderError("Error code: 500 - server error", model=self.name)
|
||
|
||
fallback = FakeModel("fallback")
|
||
wrapper = llm_factory.FallbackChatModel(primary=Boom("primary"), fallback=fallback, model="m")
|
||
with pytest.raises(ModelProviderError, match="500"):
|
||
await wrapper.ainvoke([], None)
|
||
assert fallback.calls == 0 # 服务端 500 不该消耗兜底 Key
|
||
|
||
def test_it_looks_like_a_chat_model_to_browser_use(self, clean_keys):
|
||
wrapper, _, _ = self._wrap(primary_fails=False)
|
||
assert wrapper.provider == "deepseek"
|
||
assert wrapper.name == "primary"
|
||
assert wrapper.model_name == "deepseek-chat"
|
||
assert wrapper._verified_api_keys is False
|
||
|
||
|
||
class TestLlmConfig:
|
||
def test_config_reports_both_keys_and_which_one_is_active(self, clean_keys, monkeypatch):
|
||
monkeypatch.setattr(llm_factory, "_setting", lambda key, env_key, default="": {
|
||
"llm_provider": "deepseek", "llm_model": "deepseek-chat",
|
||
"llm_base_url": "https://api.deepseek.com/v1", "llm_api_key": "PRIMARY",
|
||
}.get(key, default if isinstance(default, str) else ""))
|
||
|
||
cfg = llm_factory.get_llm_config()
|
||
assert cfg["primary_api_key"] == "PRIMARY" and cfg["fallback_api_key"] == "FALLBACK"
|
||
assert cfg["api_key"] == "PRIMARY" and cfg["using_fallback"] is False
|
||
|
||
llm_keys.mark_primary_exhausted("402")
|
||
cfg = llm_factory.get_llm_config()
|
||
assert cfg["api_key"] == "FALLBACK" and cfg["using_fallback"] is True
|