""" NVIDIA Provider Tests - ADR-036 =============================== 測試 Nemotron Tool Calling 整合 注意: 這些是單元測試，不需要真實的 NVIDIA API Key """ import json 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 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 get_ai_router, AIProvider, 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 get_ai_router, AIProvider, 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()