test_key_rotation_targeted.py

#!/usr/bin/env python3
"""
Targeted Key Rotation Test - Tests key rotation with rapid successive calls
"""

import requests
import json
import time
import threading

def rapid_fire_test():
    """Send rapid requests to trigger rate limits and key rotation"""
    print("🔥 RAPID FIRE KEY ROTATION TEST")
    print("=" * 50)
    print("🎯 Sending 10 rapid requests to trigger rate limits...")
    
    # Small document with minimal questions to process quickly
    test_payload = {
        "documents": "https://hackrx-pdf-documents.s3.amazonaws.com/documents/Family_20Medicare_20Policy_20(UIN-_20UIIHLIP22070V042122)_201.pdf",
        "questions": ["What is this document?"]
    }
    
    results = []
    
    def send_request(request_id):
        try:
            print(f"🚀 Sending request #{request_id}")
            start_time = time.time()
            
            response = requests.post(
                "http://127.0.0.1:8000/hackrx/run",
                json=test_payload,
                timeout=30
            )
            
            response_time = time.time() - start_time
            
            if response.status_code == 200:
                print(f"✅ Request #{request_id} - Success in {response_time:.2f}s")
                results.append(True)
            else:
                print(f"❌ Request #{request_id} - Failed: {response.status_code}")
                results.append(False)
                
        except Exception as e:
            print(f"❌ Request #{request_id} - Error: {e}")
            results.append(False)
    
    # Send requests in rapid succession
    threads = []
    for i in range(1, 11):  # 10 requests
        thread = threading.Thread(target=send_request, args=(i,))
        threads.append(thread)
        thread.start()
        time.sleep(0.1)  # 100ms between starts
    
    # Wait for all threads to complete
    for thread in threads:
        thread.join()
    
    successful = sum(results)
    print(f"\n📊 RAPID FIRE RESULTS")
    print(f"✅ Successful: {successful}/10")
    print(f"❌ Failed: {10-successful}/10")
    print(f"📈 Success Rate: {successful/10*100:.1f}%")

def sequential_burst_test():
    """Send a burst of sequential requests"""
    print(f"\n⚡ SEQUENTIAL BURST TEST")
    print("=" * 40)
    
    test_payload = {
        "documents": "https://hackrx-pdf-documents.s3.amazonaws.com/documents/Family_20Medicare_20Policy_20(UIN-_20UIIHLIP22070V042122)_201.pdf",
        "questions": ["What is the grace period?", "What are the benefits?"]
    }
    
    successful = 0
    
    for i in range(1, 6):  # 5 sequential requests
        try:
            print(f"📤 Sequential request #{i}...")
            start_time = time.time()
            
            response = requests.post(
                "http://127.0.0.1:8000/hackrx/run",
                json=test_payload,
                timeout=45
            )
            
            response_time = time.time() - start_time
            
            if response.status_code == 200:
                print(f"✅ Request #{i} - Success in {response_time:.2f}s")
                successful += 1
            else:
                print(f"❌ Request #{i} - Failed: {response.status_code}")
                print(f"   Error: {response.text[:150]}...")
                
        except Exception as e:
            print(f"❌ Request #{i} - Error: {e}")
        
        # Very short delay
        time.sleep(0.5)
    
    print(f"\n📊 Sequential Results: {successful}/5 successful")

if __name__ == "__main__":
    print("🔑 TARGETED KEY ROTATION TESTING")
    print("🎯 This test is designed to trigger rate limits quickly")
    print("💡 Watch the server logs for key rotation messages!")
    print()
    
    # Test 1: Rapid concurrent requests
    rapid_fire_test()
    
    # Small delay between tests
    print("\n⏳ Waiting 10 seconds before sequential test...")
    time.sleep(10)
    
    # Test 2: Sequential burst
    sequential_burst_test()
    
    print(f"\n✨ Key rotation testing complete!")
    print("📋 Check server logs for detailed rotation information")