Files
awoooi/apps/api/tests/test_nvidia_provider.py
Your Name e208798531
All checks were successful
CD Pipeline / tests (push) Successful in 54s
Code Review / ai-code-review (push) Successful in 14s
CD Pipeline / build-and-deploy (push) Successful in 3m25s
CD Pipeline / post-deploy-checks (push) Successful in 1m50s
fix(ai): keep GCP Ollama lane on safe models
2026-05-05 23:37:33 +08:00

663 lines
21 KiB
Python

"""
NVIDIA Provider Tests - ADR-036
===============================
測試 Nemotron Tool Calling 整合
注意: 這些是單元測試,不需要真實的 NVIDIA API Key
"""
import pytest
from src.models.nvidia import (
NvidiaChoice,
NvidiaMessage,
NvidiaProviderResult,
NvidiaResponse,
NvidiaUsage,
ToolCall,
ToolCallValidationResult,
ToolDefinition,
ToolFunction,
)
from src.services.nvidia_provider import (
HIGH_RISK_TOOLS,
NvidiaProvider,
OllamaToolProvider,
create_tool_definition,
get_nvidia_provider,
reset_nvidia_provider,
)
class TestNvidiaModels:
"""測試 NVIDIA Pydantic Models"""
def test_tool_function_model(self):
"""測試 ToolFunction 模型"""
func = ToolFunction(
name="restart_pod",
arguments='{"pod_name": "api-server", "namespace": "default"}',
)
assert func.name == "restart_pod"
assert '"pod_name"' in func.arguments
def test_tool_call_model(self):
"""測試 ToolCall 模型"""
tc = ToolCall(
id="call_123",
type="function",
function=ToolFunction(
name="scale_deployment",
arguments='{"replicas": 3}',
),
)
assert tc.id == "call_123"
assert tc.function.name == "scale_deployment"
def test_nvidia_response_model(self):
"""測試 NvidiaResponse 模型"""
response = NvidiaResponse(
id="resp_123",
created=1234567890,
model="nvidia/llama-3.1-nemotron-70b-instruct",
choices=[
NvidiaChoice(
index=0,
message=NvidiaMessage(
role="assistant",
content=None,
tool_calls=[
ToolCall(
id="call_1",
function=ToolFunction(
name="restart_pod",
arguments='{"pod": "api"}',
),
)
],
),
finish_reason="tool_calls",
)
],
usage=NvidiaUsage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150,
),
)
assert response.id == "resp_123"
assert len(response.choices) == 1
assert response.choices[0].message.tool_calls is not None
assert len(response.choices[0].message.tool_calls) == 1
assert response.usage.total_tokens == 150
def test_tool_call_validation_result(self):
"""測試 ToolCallValidationResult 模型"""
result = ToolCallValidationResult(
valid=True,
tool_name="restart_pod",
arguments={"pod_name": "api", "namespace": "default"},
)
assert result.valid
assert result.tool_name == "restart_pod"
assert result.arguments["namespace"] == "default"
def test_tool_definition(self):
"""測試 ToolDefinition 模型"""
definition = ToolDefinition(
type="function",
function={
"name": "restart_pod",
"description": "Restart a Kubernetes pod",
"parameters": {
"type": "object",
"properties": {
"pod_name": {"type": "string"},
"namespace": {"type": "string"},
},
"required": ["pod_name"],
},
},
)
assert definition.type == "function"
assert definition.function["name"] == "restart_pod"
class TestNvidiaProvider:
"""測試 NvidiaProvider 類別"""
def test_singleton(self):
"""測試單例模式"""
reset_nvidia_provider()
p1 = get_nvidia_provider()
p2 = get_nvidia_provider()
assert p1 is p2
reset_nvidia_provider()
def test_high_risk_tool_detection(self):
"""測試高風險 Tool 檢測"""
provider = NvidiaProvider()
# 高風險操作
assert provider.is_high_risk_tool("delete_pod")
assert provider.is_high_risk_tool("DELETE_POD") # 大寫也應該匹配
assert provider.is_high_risk_tool("delete_deployment")
assert provider.is_high_risk_tool("scale_to_zero")
assert provider.is_high_risk_tool("drain_node")
# 非高風險操作
assert not provider.is_high_risk_tool("restart_pod")
assert not provider.is_high_risk_tool("scale_deployment")
assert not provider.is_high_risk_tool("get_logs")
def test_filter_high_risk_tools(self):
"""測試過濾高風險 Tool Calls"""
provider = NvidiaProvider()
tool_calls = [
ToolCallValidationResult(
valid=True,
tool_name="restart_pod",
arguments={"pod": "api"},
),
ToolCallValidationResult(
valid=True,
tool_name="delete_pod",
arguments={"pod": "test"},
),
ToolCallValidationResult(
valid=False,
tool_name="invalid_tool",
error="Parse error",
),
]
high_risk = provider.get_high_risk_tools(tool_calls)
assert len(high_risk) == 1
assert high_risk[0].tool_name == "delete_pod"
def test_validate_tool_calls(self):
"""測試 Tool Call 驗證"""
provider = NvidiaProvider()
# 建立模擬回應
response = NvidiaResponse(
id="resp_123",
created=1234567890,
model="nvidia/llama-3.1-nemotron-70b-instruct",
choices=[
NvidiaChoice(
index=0,
message=NvidiaMessage(
role="assistant",
tool_calls=[
ToolCall(
id="call_1",
function=ToolFunction(
name="restart_pod",
arguments='{"pod_name": "api", "namespace": "default"}',
),
),
ToolCall(
id="call_2",
function=ToolFunction(
name="invalid_tool",
arguments="not valid json{", # 無效 JSON
),
),
],
),
)
],
)
results = provider._validate_tool_calls(response)
assert len(results) == 2
assert results[0].valid
assert results[0].tool_name == "restart_pod"
assert results[0].arguments["pod_name"] == "api"
assert not results[1].valid
assert "JSON" in results[1].error
class TestCreateToolDefinition:
"""測試 Tool 定義建立函數"""
def test_create_tool_definition(self):
"""測試建立 Tool 定義"""
definition = create_tool_definition(
name="scale_deployment",
description="Scale a Kubernetes deployment",
parameters={
"type": "object",
"properties": {
"deployment": {"type": "string"},
"replicas": {"type": "integer"},
},
"required": ["deployment", "replicas"],
},
)
assert definition.type == "function"
assert definition.function["name"] == "scale_deployment"
assert definition.function["description"] == "Scale a Kubernetes deployment"
assert "replicas" in definition.function["parameters"]["properties"]
class TestHighRiskTools:
"""測試高風險 Tool 清單"""
def test_high_risk_tools_list(self):
"""確認高風險 Tool 清單包含所有必要操作"""
assert "delete_pod" in HIGH_RISK_TOOLS
assert "delete_deployment" in HIGH_RISK_TOOLS
assert "delete_namespace" in HIGH_RISK_TOOLS
assert "scale_to_zero" in HIGH_RISK_TOOLS
assert "drain_node" in HIGH_RISK_TOOLS
assert "cordon_node" in HIGH_RISK_TOOLS
def test_restart_not_high_risk(self):
"""確認 restart 不在高風險清單"""
assert "restart_pod" not in HIGH_RISK_TOOLS
assert "restart_deployment" not in HIGH_RISK_TOOLS
class TestProtocolCompliance:
"""測試 Protocol 合規性 (P2-1)"""
def test_nvidia_provider_implements_protocol(self):
"""測試 NvidiaProvider 實作 INvidiaProvider Protocol"""
from src.services.nvidia_provider import INvidiaProvider, NvidiaProvider
provider = NvidiaProvider()
assert isinstance(provider, INvidiaProvider)
def test_protocol_method_signatures(self):
"""測試 Protocol 方法簽名"""
from src.services.nvidia_provider import INvidiaProvider
# Protocol 應該定義這些方法
assert hasattr(INvidiaProvider, "tool_call")
assert hasattr(INvidiaProvider, "is_high_risk_tool")
assert hasattr(INvidiaProvider, "get_high_risk_tools")
assert hasattr(INvidiaProvider, "close")
class TestOllamaToolProviderRouting:
"""Ollama tool-calling must not pollute the GCP alert lane."""
def test_base_url_uses_hermes_resolver_lane(self, monkeypatch):
from src.services import nvidia_provider as nvidia_provider_module
captured_workloads = []
def fake_resolve(workload):
captured_workloads.append(workload)
return "http://local-111:11434"
monkeypatch.setattr(nvidia_provider_module, "resolve_ollama_endpoint", fake_resolve)
provider = OllamaToolProvider()
assert provider._base_url() == "http://local-111:11434"
assert captured_workloads == ["hermes"]
@pytest.mark.asyncio
async def test_health_check_uses_local_hermes_lane(self, monkeypatch):
from src.services import nvidia_provider as nvidia_provider_module
class _FakeResponse:
status_code = 200
class _FakeClient:
def __init__(self):
self.checked_urls = []
async def get(self, url, **kwargs):
self.checked_urls.append(url)
return _FakeResponse()
monkeypatch.setattr(
nvidia_provider_module,
"resolve_ollama_endpoint",
lambda workload: "http://local-111:11434",
)
provider = OllamaToolProvider()
client = _FakeClient()
async def _get_client():
return client
monkeypatch.setattr(provider, "_get_client", _get_client)
assert await provider.health_check() is True
assert client.checked_urls == ["http://local-111:11434/api/tags"]
class TestEdgeCases:
"""邊界測試案例 (P2-2)"""
@pytest.mark.asyncio
async def test_api_key_not_set(self):
"""測試 API Key 未設定時返回 fallback"""
provider = NvidiaProvider(api_key="") # 明確設定空 key
result = await provider.tool_call(
messages=[{"role": "user", "content": "test"}],
tools=[],
)
assert not result.success
assert result.fallback_triggered
assert "NVIDIA_API_KEY" in result.error
def test_empty_tool_calls_response(self):
"""測試無 Tool Call 的回應"""
provider = NvidiaProvider()
# 建立沒有 tool_calls 的回應
response = NvidiaResponse(
id="resp_123",
created=1234567890,
model="nvidia/nemotron-mini-4b-instruct",
choices=[
NvidiaChoice(
index=0,
message=NvidiaMessage(
role="assistant",
content="I cannot help with that.",
tool_calls=None, # 無 tool_calls
),
)
],
)
results = provider._validate_tool_calls(response)
assert len(results) == 0
def test_empty_choices_response(self):
"""測試空 choices 的回應"""
provider = NvidiaProvider()
response = NvidiaResponse(
id="resp_123",
created=1234567890,
model="nvidia/nemotron-mini-4b-instruct",
choices=[], # 空 choices
)
results = provider._validate_tool_calls(response)
assert len(results) == 0
def test_provider_result_model(self):
"""測試 NvidiaProviderResult 模型各種狀態"""
# 成功結果
success_result = NvidiaProviderResult(
success=True,
tool_calls=[
ToolCallValidationResult(
valid=True,
tool_name="restart_pod",
arguments={"pod": "api"},
)
],
usage=NvidiaUsage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150,
),
latency_ms=1000.0,
)
assert success_result.success
assert len(success_result.tool_calls) == 1
assert success_result.usage.total_tokens == 150
# 失敗結果
fail_result = NvidiaProviderResult(
success=False,
error="Connection timeout",
fallback_triggered=True,
)
assert not fail_result.success
assert fail_result.fallback_triggered
assert "timeout" in fail_result.error.lower()
def test_all_high_risk_tools_covered(self):
"""確保所有危險操作都被標記為高風險"""
dangerous_operations = [
"delete_pod",
"delete_deployment",
"delete_namespace",
"delete_service",
"delete_configmap",
"delete_secret",
"scale_to_zero",
"drain_node",
"cordon_node",
"delete_pvc",
"delete_pv",
]
for op in dangerous_operations:
assert op in HIGH_RISK_TOOLS, f"{op} should be in HIGH_RISK_TOOLS"
class TestAIRouterNvidiaIntegration:
"""測試 AIRouter NVIDIA 整合"""
def test_nvidia_provider_in_router(self):
"""測試 AIProviderEnum 包含 NEMOTRON (Phase 24 B3: NVIDIA → NEMOTRON)"""
from src.services.ai_router import AIProviderEnum
assert hasattr(AIProviderEnum, "NEMOTRON")
assert AIProviderEnum.NEMOTRON.value == "nemotron"
def test_tool_calling_route(self):
"""測試 Tool Calling 路由"""
from src.services.ai_router import (
AIProviderEnum,
get_ai_router,
reset_ai_router,
)
reset_ai_router()
router = get_ai_router()
provider, model, fallback_chain = router.route_tool_calling()
assert provider == AIProviderEnum.NEMOTRON
assert "nvidia" in model.lower() or "nemotron" in model.lower()
# Fallback 應該包含 Gemini 和 Claude
fallback_providers = [p for p, _ in fallback_chain]
assert AIProviderEnum.GEMINI in fallback_providers
assert AIProviderEnum.CLAUDE in fallback_providers
reset_ai_router()
def test_existing_routing_not_affected(self):
"""測試現有路由規則不受影響"""
from src.services.ai_router import (
AIProviderEnum,
get_ai_router,
reset_ai_router,
)
reset_ai_router()
router = get_ai_router()
# 測試同步路由 (不涉及 NEMOTRON)
decision = router.route_sync("重啟 api pod")
# 應該還是使用 Ollama (低複雜度)
assert decision.selected_provider in [
AIProviderEnum.OLLAMA,
AIProviderEnum.GEMINI,
AIProviderEnum.CLAUDE,
]
# NEMOTRON 不應該出現在一般路由中
assert decision.selected_provider != AIProviderEnum.NEMOTRON
reset_ai_router()
class TestRateLimiterIntegration:
"""測試 Rate Limiter 整合 (P2-2)"""
def test_nvidia_in_rate_limits(self):
"""測試 NVIDIA 在 Rate Limits 配置中"""
from src.services.ai_rate_limiter import RATE_LIMITS
assert "openclaw_nemo" in RATE_LIMITS
assert "rpm" in RATE_LIMITS["openclaw_nemo"]
assert "daily_requests" in RATE_LIMITS["openclaw_nemo"]
def test_nvidia_rate_limit_values(self):
"""測試 NVIDIA Rate Limit 值正確"""
from src.services.ai_rate_limiter import RATE_LIMITS
nvidia_limits = RATE_LIMITS["openclaw_nemo"]
assert nvidia_limits["rpm"] == 10 # 2026-04-01 ogt: 放寬到 10 (原 5)
assert nvidia_limits["daily_requests"] == 99999 # 免費版無限制,設大數 (原 100)
assert nvidia_limits["daily_tokens"] == 9999999 # 免費版無限制 (原 50_000)
def test_nvidia_in_cost_limits(self):
"""測試 NVIDIA 在成本限制中 (免費 tier)"""
from src.services.ai_rate_limiter import COST_LIMITS
assert "openclaw_nemo" in COST_LIMITS
assert COST_LIMITS["openclaw_nemo"]["total_cost_usd"] == 999999.0 # 免費 Tier 無成本限制 (2026-04-01 ogt: 原 0.0)
class TestCircuitBreaker:
"""P3-1: Circuit Breaker 測試"""
def test_circuit_breaker_initial_state(self):
"""測試 Circuit Breaker 初始狀態"""
from src.services.nvidia_provider import CircuitBreaker, CircuitState
cb = CircuitBreaker()
assert cb.state == CircuitState.CLOSED
assert cb.can_execute()
def test_circuit_breaker_opens_after_failures(self):
"""測試連續失敗後斷路"""
from src.services.nvidia_provider import CircuitBreaker, CircuitState
cb = CircuitBreaker(failure_threshold=3)
# 連續 3 次失敗
cb.record_failure()
assert cb.state == CircuitState.CLOSED
cb.record_failure()
assert cb.state == CircuitState.CLOSED
cb.record_failure()
assert cb.state == CircuitState.OPEN
assert not cb.can_execute()
def test_circuit_breaker_success_resets_count(self):
"""測試成功重置失敗計數"""
from src.services.nvidia_provider import CircuitBreaker, CircuitState
cb = CircuitBreaker(failure_threshold=3)
cb.record_failure()
cb.record_failure()
cb.record_success() # 重置
# 需要再 3 次失敗才能斷路
cb.record_failure()
cb.record_failure()
assert cb.state == CircuitState.CLOSED
cb.record_failure()
assert cb.state == CircuitState.OPEN
def test_circuit_breaker_half_open_recovery(self):
"""測試半開狀態恢復"""
from src.services.nvidia_provider import CircuitBreaker, CircuitState
cb = CircuitBreaker(failure_threshold=1, recovery_timeout=0.1)
cb.record_failure() # 觸發斷路
assert cb.state == CircuitState.OPEN
import time
time.sleep(0.15) # 等待恢復
# 檢查狀態會觸發 HALF_OPEN 轉換
assert cb.state == CircuitState.HALF_OPEN
assert cb.can_execute()
# 成功後回到 CLOSED
cb.record_success()
assert cb.state == CircuitState.CLOSED
def test_circuit_breaker_half_open_failure(self):
"""測試半開狀態失敗重新斷路"""
from src.services.nvidia_provider import CircuitBreaker, CircuitState
cb = CircuitBreaker(failure_threshold=1, recovery_timeout=0.1)
cb.record_failure()
assert cb.state == CircuitState.OPEN
import time
time.sleep(0.15)
assert cb.state == CircuitState.HALF_OPEN
# 失敗,重新斷路
cb.record_failure()
assert cb.state == CircuitState.OPEN
def test_provider_has_circuit_breaker(self):
"""測試 NvidiaProvider 有 Circuit Breaker"""
from src.services.nvidia_provider import NvidiaProvider
provider = NvidiaProvider()
assert hasattr(provider, "_circuit_breaker")
class TestPrometheusMetrics:
"""P3-3: Prometheus Metrics 測試"""
def test_metrics_defined(self):
"""測試 Prometheus Metrics 已定義"""
from src.services.nvidia_provider import (
NVIDIA_CIRCUIT_BREAKER_STATE,
NVIDIA_LATENCY_HISTOGRAM,
NVIDIA_REQUESTS_TOTAL,
)
assert NVIDIA_REQUESTS_TOTAL is not None
assert NVIDIA_LATENCY_HISTOGRAM is not None
assert NVIDIA_CIRCUIT_BREAKER_STATE is not None
def test_metrics_labels(self):
"""測試 Metrics Labels 正確"""
from src.services.nvidia_provider import NVIDIA_REQUESTS_TOTAL
# 應該能建立帶 label 的 metric
metric = NVIDIA_REQUESTS_TOTAL.labels(status="test", tool_name="test_tool")
assert metric is not None
class TestExponentialBackoff:
"""P3-2: 指數退避測試"""
def test_backoff_constants_defined(self):
"""測試退避常數已定義"""
from src.services.nvidia_provider import (
RETRY_BASE_DELAY,
RETRY_EXPONENTIAL_BASE,
RETRY_MAX_DELAY,
)
assert RETRY_BASE_DELAY == 1.0
assert RETRY_MAX_DELAY == 30.0
assert RETRY_EXPONENTIAL_BASE == 2