147 lines
5.2 KiB
Python
147 lines
5.2 KiB
Python
import asyncio
|
|
from unittest.mock import AsyncMock, patch, MagicMock
|
|
from db.models import AIProviderProfile
|
|
from core.llm import LLMClient
|
|
from services.admin_bot import get_persian_main_menu, AdminBotService
|
|
|
|
|
|
def test_main_menu_pause_toggle():
|
|
menu_running = get_persian_main_menu(is_paused=False)
|
|
button_texts_running = [getattr(b.button, 'text', '') for row in menu_running for b in row]
|
|
assert "🛑 توقف اضطراری سیستم" in button_texts_running
|
|
assert "▶️ راهاندازی و ادامه سیستم" not in button_texts_running
|
|
|
|
menu_paused = get_persian_main_menu(is_paused=True)
|
|
button_texts_paused = [getattr(b.button, 'text', '') for row in menu_paused for b in row]
|
|
assert "▶️ راهاندازی و ادامه سیستم" in button_texts_paused
|
|
assert "🛑 توقف اضطراری سیستم" not in button_texts_paused
|
|
|
|
|
|
async def test_llm_fallback_chain_success_on_fallback():
|
|
# Setup 2 providers: Primary (fails) -> Fallback (succeeds)
|
|
profile_1 = AIProviderProfile(
|
|
id=1,
|
|
name="Primary AI",
|
|
provider_type="openai",
|
|
model="gpt-4o",
|
|
base_url="https://api.openai.com/v1",
|
|
api_key="key1",
|
|
is_active=True,
|
|
fallback_provider_id=2
|
|
)
|
|
profile_2 = AIProviderProfile(
|
|
id=2,
|
|
name="Backup AI",
|
|
provider_type="gemini",
|
|
model="gemini-1.5-flash",
|
|
base_url="",
|
|
api_key="key2",
|
|
is_active=False,
|
|
fallback_provider_id=None
|
|
)
|
|
|
|
repo_mock = AsyncMock()
|
|
repo_mock.get_active_provider_profile.return_value = profile_1
|
|
repo_mock.get_provider_profiles.return_value = [profile_1, profile_2]
|
|
repo_mock.record_ai_log = AsyncMock()
|
|
|
|
fallback_alerts = []
|
|
|
|
async def on_fallback(from_p, to_p, err, step):
|
|
fallback_alerts.append((from_p.name, to_p.name, err, step))
|
|
|
|
client = LLMClient(
|
|
repo=repo_mock,
|
|
on_fallback_alert=on_fallback,
|
|
)
|
|
client.max_retries_per_model = 0 # fail fast for test
|
|
|
|
# Mock _call_openai to fail and _call_gemini to succeed
|
|
with patch.object(client, "_call_openai", side_effect=RuntimeError("Primary Provider Timeout")):
|
|
with patch.object(client, "_call_gemini", return_value={"decision": "accept", "rewritten_text": "Success from backup"}):
|
|
res = await client.generate_json("Test prompt", system_prompt="Test sys")
|
|
assert res["decision"] == "accept"
|
|
assert res["rewritten_text"] == "Success from backup"
|
|
assert len(fallback_alerts) == 1
|
|
assert fallback_alerts[0][0] == "Primary AI"
|
|
assert fallback_alerts[0][1] == "Backup AI"
|
|
assert "Primary Provider Timeout" in fallback_alerts[0][2]
|
|
assert fallback_alerts[0][3] == 1
|
|
|
|
|
|
async def test_llm_fallback_chain_all_fail_alert():
|
|
profile_1 = AIProviderProfile(
|
|
id=1,
|
|
name="Primary AI",
|
|
provider_type="openai",
|
|
model="gpt-4o",
|
|
is_active=True,
|
|
fallback_provider_id=2
|
|
)
|
|
profile_2 = AIProviderProfile(
|
|
id=2,
|
|
name="Backup AI 1",
|
|
provider_type="openai",
|
|
model="gpt-4o-mini",
|
|
is_active=False,
|
|
fallback_provider_id=3
|
|
)
|
|
profile_3 = AIProviderProfile(
|
|
id=3,
|
|
name="Backup AI 2",
|
|
provider_type="gemini",
|
|
model="gemini-1.5-flash",
|
|
is_active=False,
|
|
fallback_provider_id=None
|
|
)
|
|
|
|
repo_mock = AsyncMock()
|
|
repo_mock.get_active_provider_profile.return_value = profile_1
|
|
repo_mock.get_provider_profiles.return_value = [profile_1, profile_2, profile_3]
|
|
repo_mock.record_ai_log = AsyncMock()
|
|
|
|
fallback_alerts = []
|
|
chain_failure_alerts = []
|
|
|
|
async def on_fallback(from_p, to_p, err, step):
|
|
fallback_alerts.append((from_p.name, to_p.name, err, step))
|
|
|
|
async def on_chain_failure(chain, err):
|
|
chain_failure_alerts.append((len(chain), err))
|
|
|
|
client = LLMClient(
|
|
repo=repo_mock,
|
|
on_fallback_alert=on_fallback,
|
|
on_chain_failure_alert=on_chain_failure,
|
|
)
|
|
client.max_retries_per_model = 0
|
|
|
|
with patch.object(client, "_call_openai", side_effect=RuntimeError("OpenAI Error")):
|
|
with patch.object(client, "_call_gemini", side_effect=RuntimeError("Gemini Quota Exceeded")):
|
|
try:
|
|
await client.generate_json("Test prompt")
|
|
assert False, "Should have raised exception"
|
|
except Exception as e:
|
|
assert "Gemini Quota Exceeded" in str(e) or "OpenAI Error" in str(e)
|
|
|
|
# Check fallbacks: 1 -> 2, 2 -> 3
|
|
assert len(fallback_alerts) == 2
|
|
assert fallback_alerts[0][0] == "Primary AI"
|
|
assert fallback_alerts[0][1] == "Backup AI 1"
|
|
assert fallback_alerts[1][0] == "Backup AI 1"
|
|
assert fallback_alerts[1][1] == "Backup AI 2"
|
|
|
|
# Check chain failure alert triggered
|
|
assert len(chain_failure_alerts) == 1
|
|
assert chain_failure_alerts[0][0] == 3
|
|
|
|
|
|
async def main():
|
|
test_main_menu_pause_toggle()
|
|
await test_llm_fallback_chain_success_on_fallback()
|
|
await test_llm_fallback_chain_all_fail_alert()
|
|
print("All system stop and multi-provider AI fallback chain tests passed successfully!")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|