Upload production/tests/test_fallback.py

production/tests/test_fallback.py

@@ -1,6 +1,4 @@
-"""
-Integration tests for NIM ↔ Cloudflare fallback logic.
-"""
+"""Integration tests for NIM ↔ Cloudflare ↔ Gemini fallback logic."""
 
 import asyncio
 import pytest
@@ -34,7 +32,7 @@ async def mock_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))
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", enabled=True))
         provider, config = await mgr.get_active_provider()
         assert provider == "nim"
         assert config["api_base"] == "https://integrate.api.nvidia.com/v1"
@@ -45,22 +43,42 @@ class TestFallbackManager:
             {"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))
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", 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):
+    async def test_falls_to_tertiary_when_secondary_open(self, mock_redis):
+        import production_server
+        old_gemini = production_server.GEMINI_API_KEY
+        production_server.GEMINI_API_KEY = "test-key"
+        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", tertiary="gemini", enabled=True))
+            provider, config = await mgr.get_active_provider()
+            assert provider == "gemini"
+            assert "generativelanguage" in config["api_base"]
+            assert config["api_key"] == "test-key"
+        finally:
+            production_server.GEMINI_API_KEY = old_gemini
+
+    @pytest.mark.asyncio
+    async def test_falls_to_mlx_when_all_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": "open", "failures": 5, "last_failure": 9999999999},
                 {"state": "closed", "failures": 0, "last_failure": 0},
             ])
-            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=True))
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", enabled=True))
             provider, config = await mgr.get_active_provider()
             assert provider == "mlx"
         finally:
@@ -75,8 +93,9 @@ class TestFallbackManager:
             mock_redis.get_circuit_state = AsyncMock(side_effect=[
                 {"state": "open", "failures": 5, "last_failure": 9999999999},
                 {"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))
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", enabled=True))
             with pytest.raises(HTTPException) as exc_info:
                 await mgr.get_active_provider()
             assert exc_info.value.status_code == 503
@@ -85,30 +104,28 @@ class TestFallbackManager:
 
     @pytest.mark.asyncio
     async def test_respects_disabled_fallback(self, mock_redis):
-        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", enabled=False))
+        mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", enabled=False))
         provider, config = await mgr.get_active_provider()
         assert provider == "nim"
 
     @pytest.mark.asyncio
-    async def test_cloudflare_config(self, mock_redis):
+    async def test_gemini_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"
+        old_gemini = production_server.GEMINI_API_KEY
+        production_server.GEMINI_API_KEY = "gemini-test-key"
         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))
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", 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"
+            assert provider == "gemini"
+            assert config["rpm_limit"] == 60
+            assert config["cost_per_1m_input"] == 0.075
         finally:
-            production_server.CLOUDFLARE_API_KEY = old_cf_key
-            production_server.CLOUDFLARE_ACCOUNT_ID = old_cf_id
+            production_server.GEMINI_API_KEY = old_gemini
 
     @pytest.mark.asyncio
     async def test_nim_config(self, mock_redis):
@@ -116,7 +133,7 @@ class TestFallbackManager:
         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))
+            mgr = FallbackManager(mock_redis, FallbackConfig(primary="nim", secondary="cloudflare", tertiary="gemini", enabled=True))
             provider, config = await mgr.get_active_provider()
             assert provider == "nim"
             assert config["api_base"] == "https://custom.nvidia.com/v1"