verify_fallback.py

import os
import sys
from unittest.mock import MagicMock, patch

# Mocking modules that might not be available or needed for this test
sys.modules['cv2'] = MagicMock()
sys.modules['whisper'] = MagicMock()

# Set dummy env vars BEFORE importing agent
os.environ["OPENROUTER_API_KEY"] = "dummy"
os.environ["GOOGLE_API_KEY"] = "dummy"
os.environ["GROQ_API_KEY"] = "dummy"
os.environ["NVIDIA_API_KEY"] = "dummy"
os.environ["VERCEL_API_KEY"] = "dummy"

# Add the current directory to path so we can import agent
sys.path.append(os.getcwd())

import agent
from langchain_core.messages import HumanMessage

def test_fallback_logic():
    print("Testing fallback logic...")
    
    # Mock the invoke method for each tier's model
    # Tiers: 0:OpenRouter, 1:Gemini, 2:Groq, 3:NVIDIA, 4:Vercel
    
    with patch('agent.openrouter_model.invoke') as mock_openrouter, \
         patch('agent.gemini_model.invoke') as mock_gemini, \
         patch('agent.model.invoke') as mock_groq, \
         patch('agent.nvidia_model.invoke') as mock_nvidia, \
         patch('agent.vercel_model.invoke') as mock_vercel:
        
        # Simulate failure for all tiers up to NVIDIA
        mock_openrouter.side_effect = Exception("Rate limit (429)")
        mock_gemini.side_effect = Exception("Rate limit (429)")
        mock_groq.side_effect = Exception("Rate limit (429)")
        
        # NVIDIA should succeed
        mock_nvidia.return_value = MagicMock(content="NVIDIA response")
        
        msgs = [HumanMessage(content="Hello")]
        response, tier_idx = agent.smart_invoke(msgs, use_tools=False)
        
        print(f"Response from tier {tier_idx}: {response.content}")
        assert tier_idx == 3
        assert response.content == "NVIDIA response"
        print("Fallback to NVIDIA successful!")

        # Now simulate failure up to Vercel
        mock_nvidia.side_effect = Exception("Rate limit (429)")
        mock_vercel.return_value = MagicMock(content="Vercel response")
        
        response, tier_idx = agent.smart_invoke(msgs, use_tools=False)
        print(f"Response from tier {tier_idx}: {response.content}")
        assert tier_idx == 4
        assert response.content == "Vercel response"
        print("Fallback to Vercel successful!")

if __name__ == "__main__":
    try:
        test_fallback_logic()
        print("All fallback tests passed!")
    except Exception as e:
        print(f"Test failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)