154 lines
5.3 KiB
Python
154 lines
5.3 KiB
Python
"""
|
||
P0 DIAGNOSE Privacy-First Routing Tests
|
||
========================================
|
||
測試 AIRouter local chain 隔離 + DIAGNOSE timeout 路由
|
||
|
||
建立時間: 2026-04-04 (台北時區)
|
||
建立者: Claude Code (P0 DIAGNOSE Privacy-First)
|
||
"""
|
||
|
||
import os
|
||
os.environ.setdefault("MOCK_MODE", "true")
|
||
|
||
import pytest
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
|
||
class TestNemotronPerTaskTimeout:
|
||
"""Nemotron 支援 per-task timeout"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_diagnose_uses_diagnose_timeout(self):
|
||
"""DIAGNOSE context 應使用 NEMOTRON_DIAGNOSE_TIMEOUT_SECONDS"""
|
||
from src.services.ai_providers.nemotron import NemotronProvider
|
||
|
||
provider = NemotronProvider()
|
||
|
||
# 建立 mock nvidia provider
|
||
mock_nvidia = MagicMock()
|
||
mock_result = MagicMock()
|
||
mock_result.tool_calls = []
|
||
mock_nvidia.tool_call = AsyncMock(return_value=mock_result)
|
||
|
||
with patch.object(provider, '_get_nvidia', return_value=mock_nvidia):
|
||
result = await provider.analyze(
|
||
prompt="測試診斷",
|
||
context={"task_type": "diagnose"},
|
||
)
|
||
|
||
assert result.success is True
|
||
mock_nvidia.tool_call.assert_called_once()
|
||
|
||
|
||
class TestLocalFallbackChain:
|
||
"""require_local=True 時只走 local chain，全部失敗 → REJECT，不觸碰雲端"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_require_local_skips_cloud_providers(self):
|
||
"""require_local=True 時，cloud provider 不被呼叫"""
|
||
import os
|
||
from src.services.ai_router import AIRouterExecutor, AIProviderRegistry
|
||
from src.services.ai_providers.interfaces import AIResult
|
||
|
||
registry = AIProviderRegistry()
|
||
|
||
# Mock: Ollama 成功
|
||
mock_ollama = AsyncMock()
|
||
mock_ollama.name = "ollama"
|
||
mock_ollama.privacy_level = "local"
|
||
mock_ollama.is_enabled = True
|
||
mock_ollama.capabilities = {"rca", "chat"}
|
||
mock_ollama.analyze = AsyncMock(return_value=AIResult(
|
||
raw_response="本地診斷結果",
|
||
success=True,
|
||
provider="ollama",
|
||
))
|
||
mock_ollama.health_check = AsyncMock(return_value=True)
|
||
|
||
# Mock: Gemini（不應該被呼叫）
|
||
mock_gemini = AsyncMock()
|
||
mock_gemini.name = "gemini"
|
||
mock_gemini.privacy_level = "cloud"
|
||
mock_gemini.is_enabled = True
|
||
mock_gemini.analyze = AsyncMock(return_value=AIResult(
|
||
raw_response="雲端結果",
|
||
success=True,
|
||
provider="gemini",
|
||
))
|
||
|
||
registry._providers = {
|
||
"ollama": mock_ollama,
|
||
"gemini": mock_gemini,
|
||
}
|
||
|
||
executor = AIRouterExecutor(registry)
|
||
|
||
# 暫時關閉 MOCK_MODE，測試真實執行路徑
|
||
with patch("src.services.ai_router._settings") as mock_settings:
|
||
mock_settings.MOCK_MODE = False
|
||
result = await executor.execute(
|
||
prompt="診斷這個問題",
|
||
provider_order=["ollama", "gemini"],
|
||
require_local=True,
|
||
)
|
||
|
||
assert result.success is True
|
||
assert result.provider == "ollama"
|
||
mock_gemini.analyze.assert_not_called()
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_require_local_all_fail_returns_reject(self):
|
||
"""require_local=True 且所有 local provider 失敗 → 回傳明確錯誤"""
|
||
import os
|
||
from src.services.ai_router import AIRouterExecutor, AIProviderRegistry
|
||
from src.services.ai_providers.interfaces import AIResult
|
||
|
||
registry = AIProviderRegistry()
|
||
|
||
# Mock: Ollama 失敗
|
||
mock_ollama = AsyncMock()
|
||
mock_ollama.name = "ollama"
|
||
mock_ollama.privacy_level = "local"
|
||
mock_ollama.is_enabled = True
|
||
mock_ollama.capabilities = {"rca", "chat"}
|
||
mock_ollama.analyze = AsyncMock(return_value=AIResult(
|
||
raw_response="",
|
||
success=False,
|
||
provider="ollama",
|
||
error="timeout",
|
||
))
|
||
mock_ollama.health_check = AsyncMock(return_value=False)
|
||
|
||
registry._providers = {
|
||
"ollama": mock_ollama,
|
||
}
|
||
|
||
executor = AIRouterExecutor(registry)
|
||
|
||
# 暫時關閉 MOCK_MODE + 讓 telegram import 失敗（不影響主流程）
|
||
with patch("src.services.ai_router._settings") as mock_settings:
|
||
mock_settings.MOCK_MODE = False
|
||
result = await executor.execute(
|
||
prompt="診斷這個問題",
|
||
provider_order=["ollama"],
|
||
require_local=True,
|
||
)
|
||
|
||
assert result.success is False
|
||
assert result.error == "local_providers_unavailable"
|
||
|
||
|
||
class TestDiagnoseIntentOverride:
|
||
"""DIAGNOSE intent 應優先路由至 Nemotron（非 FORCE_LOCAL 情境）"""
|
||
|
||
def test_diagnose_override_is_nemotron(self):
|
||
"""_intent_provider_overrides[DIAGNOSE] 應為 NEMOTRON"""
|
||
from src.services.ai_router import AIRouter, AIProviderEnum
|
||
from src.services.intent_classifier import IntentType
|
||
|
||
router = AIRouter()
|
||
override = router._intent_provider_overrides.get(IntentType.DIAGNOSE)
|
||
assert override == AIProviderEnum.NEMOTRON, (
|
||
f"DIAGNOSE 應路由至 NEMOTRON，實際為 {override}"
|
||
)
|