cd6da9c8d6
ai_rate_limiter.py 在 2026-03-31 更新了 NVIDIA 免費版限制值， 但測試未同步更新導致失敗: - rpm: 5 → 10 (放寬並發控制) - daily_requests: 100 → 99999 (免費版無限制) - daily_tokens: 50_000 → 9999999 (免費版無限制) - total_cost_usd: 0.0 → 999999.0 (修復 $0>=0 永遠 True bug) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
602 lines
19 KiB
Python
602 lines
19 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,
|
|
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 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):
|
|
"""測試 AIProvider 包含 NVIDIA"""
|
|
from src.services.ai_router import AIProvider
|
|
|
|
assert hasattr(AIProvider, "NVIDIA")
|
|
assert AIProvider.NVIDIA.value == "nvidia"
|
|
|
|
def test_tool_calling_route(self):
|
|
"""測試 Tool Calling 路由"""
|
|
from src.services.ai_router import AIProvider, get_ai_router, reset_ai_router
|
|
|
|
reset_ai_router()
|
|
router = get_ai_router()
|
|
|
|
provider, model, fallback_chain = router.route_tool_calling()
|
|
|
|
assert provider == AIProvider.NVIDIA
|
|
assert "nvidia" in model.lower() or "nemotron" in model.lower()
|
|
# Fallback 應該包含 Gemini 和 Claude
|
|
fallback_providers = [p for p, _ in fallback_chain]
|
|
assert AIProvider.GEMINI in fallback_providers
|
|
assert AIProvider.CLAUDE in fallback_providers
|
|
|
|
reset_ai_router()
|
|
|
|
def test_existing_routing_not_affected(self):
|
|
"""測試現有路由規則不受影響"""
|
|
from src.services.ai_router import AIProvider, get_ai_router, reset_ai_router
|
|
|
|
reset_ai_router()
|
|
router = get_ai_router()
|
|
|
|
# 測試同步路由 (不涉及 NVIDIA)
|
|
decision = router.route_sync("重啟 api pod")
|
|
|
|
# 應該還是使用 Ollama (低複雜度)
|
|
assert decision.selected_provider in [
|
|
AIProvider.OLLAMA,
|
|
AIProvider.GEMINI,
|
|
AIProvider.CLAUDE,
|
|
]
|
|
# NVIDIA 不應該出現在一般路由中
|
|
assert decision.selected_provider != AIProvider.NVIDIA
|
|
|
|
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 "nvidia" in RATE_LIMITS
|
|
assert "rpm" in RATE_LIMITS["nvidia"]
|
|
assert "daily_requests" in RATE_LIMITS["nvidia"]
|
|
|
|
def test_nvidia_rate_limit_values(self):
|
|
"""測試 NVIDIA Rate Limit 值正確"""
|
|
from src.services.ai_rate_limiter import RATE_LIMITS
|
|
|
|
nvidia_limits = RATE_LIMITS["nvidia"]
|
|
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 "nvidia" in COST_LIMITS
|
|
assert COST_LIMITS["nvidia"]["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
|