awoooi/apps/api/tests/test_nvidia_provider.py

"""
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):
        """測試 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