Upload production/tests/test_fallback.py

production/tests/test_fallback.py

@@ -0,0 +1,150 @@
+"""
+Integration tests for NIM ↔ Cloudflare fallback logic.
+"""
+
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import sys
+import os
+import time
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+from production_server import (
+    FallbackManager,
+    FallbackConfig,
+    CircuitBreaker,
+    RedisManager,
+    HTTPException,
+)
+
+
+@pytest.fixture
+async def mock_redis():
+    redis = MagicMock(spec=RedisManager)
+    redis.get_circuit_state = AsyncMock(return_value={"state": "closed", "failures": 0, "last_failure": 0})
+    redis.set_circuit_state = AsyncMock()
+    redis.get_cache = AsyncMock(return_value=None)
+    redis.set_cache = AsyncMock()
+    redis.check_rate_limit = AsyncMock(return_value=(True, 0.0))
+    return redis
+
+
+class TestFallbackManager:
+    @pytest.mark.asyncio
+    async def test_uses_primary_when_healthy(self, mock_redis):
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+        provider, config = await mgr.get_active_provider()
+        assert provider == "nim"
+        assert config["api_base"] == "https://integrate.api.nvidia.com/v1"
+
+    @pytest.mark.asyncio
+    async def test_falls_back_when_primary_open(self, mock_redis):
+        mock_redis.get_circuit_state = AsyncMock(side_effect=[
+            {"state": "open", "failures": 5, "last_failure": 9999999999},
+            {"state": "closed", "failures": 0, "last_failure": 0},
+        ])
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+        provider, config = await mgr.get_active_provider()
+        assert provider == "cloudflare"
+
+    @pytest.mark.asyncio
+    async def test_falls_to_mlx_when_both_cloud_down(self, mock_redis):
+        import production_server
+        old_mlx = production_server.MLX_ENABLED
+        production_server.MLX_ENABLED = True
+        try:
+            mock_redis.get_circuit_state = AsyncMock(side_effect=[
+                {"state": "open", "failures": 5, "last_failure": 9999999999},
+                {"state": "open", "failures": 5, "last_failure": 9999999999},
+                {"state": "closed", "failures": 0, "last_failure": 0},
+            ])
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+            provider, config = await mgr.get_active_provider()
+            assert provider == "mlx"
+        finally:
+            production_server.MLX_ENABLED = old_mlx
+
+    @pytest.mark.asyncio
+    async def test_raises_when_all_down(self, mock_redis):
+        import production_server
+        old_mlx = production_server.MLX_ENABLED
+        production_server.MLX_ENABLED = False
+        try:
+            mock_redis.get_circuit_state = AsyncMock(side_effect=[
+                {"state": "open", "failures": 5, "last_failure": 9999999999},
+                {"state": "open", "failures": 5, "last_failure": 9999999999},
+            ])
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+            with pytest.raises(HTTPException) as exc_info:
+                await mgr.get_active_provider()
+            assert exc_info.value.status_code == 503
+        finally:
+            production_server.MLX_ENABLED = old_mlx
+
+    @pytest.mark.asyncio
+    async def test_respects_disabled_fallback(self, mock_redis):
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=False))
+        provider, config = await mgr.get_active_provider()
+        assert provider == "nim"
+
+    @pytest.mark.asyncio
+    async def test_cloudflare_config(self, mock_redis):
+        import production_server
+        old_cf_key = production_server.CLOUDFLARE_API_KEY
+        old_cf_id = production_server.CLOUDFLARE_ACCOUNT_ID
+        production_server.CLOUDFLARE_API_KEY = "test-key"
+        production_server.CLOUDFLARE_ACCOUNT_ID = "test-account"
+        try:
+            mock_redis.get_circuit_state = AsyncMock(side_effect=[
+                {"state": "open", "failures": 5, "last_failure": 9999999999},
+                {"state": "closed", "failures": 0, "last_failure": 0},
+            ])
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+            provider, config = await mgr.get_active_provider()
+            assert provider == "cloudflare"
+            assert "api.cloudflare.com" in config["api_base"]
+            assert config["api_key"] == "test-key"
+        finally:
+            production_server.CLOUDFLARE_API_KEY = old_cf_key
+            production_server.CLOUDFLARE_ACCOUNT_ID = old_cf_id
+
+    @pytest.mark.asyncio
+    async def test_nim_config(self, mock_redis):
+        import production_server
+        old_nim = production_server.NIM_API_BASE
+        production_server.NIM_API_BASE = "https://custom.nvidia.com/v1"
+        try:
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+            provider, config = await mgr.get_active_provider()
+            assert provider == "nim"
+            assert config["api_base"] == "https://custom.nvidia.com/v1"
+        finally:
+            production_server.NIM_API_BASE = old_nim
+
+
+class TestCircuitBreakerFallback:
+    @pytest.mark.asyncio
+    async def test_circuit_open_then_half_open_then_closed(self, mock_redis):
+        cb = CircuitBreaker(mock_redis, "nim")
+        assert await cb.can_execute()
+        for _ in range(5):
+            await cb.record_failure()
+        mock_redis.get_circuit_state = AsyncMock(return_value={
+            "state": "open", "failures": 5, "last_failure": time.time()
+        })
+        assert not await cb.can_execute()
+        mock_redis.get_circuit_state = AsyncMock(return_value={
+            "state": "open", "failures": 5, "last_failure": time.time() - 120
+        })
+        assert await cb.can_execute()
+        await cb.record_success()
+        mock_redis.get_circuit_state = AsyncMock(return_value={
+            "state": "closed", "failures": 0, "last_failure": 0
+        })
+        assert await cb.can_execute()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])