talk2me/test_size_limits.py

#!/usr/bin/env python3
"""
Test script for request size limits
"""
import requests
import json
import io
import os

BASE_URL = "http://localhost:5005"

def test_json_size_limit():
    """Test JSON payload size limit"""
    print("\n=== Testing JSON Size Limit ===")
    
    # Create a large JSON payload (over 1MB)
    large_data = {
        "text": "x" * (2 * 1024 * 1024),  # 2MB of text
        "source_lang": "English",
        "target_lang": "Spanish"
    }
    
    try:
        response = requests.post(f"{BASE_URL}/translate", json=large_data)
        print(f"Status: {response.status_code}")
        if response.status_code == 413:
            print(f"✓ Correctly rejected large JSON: {response.json()}")
        else:
            print(f"✗ Should have rejected large JSON")
    except Exception as e:
        print(f"Error: {e}")

def test_audio_size_limit():
    """Test audio file size limit"""
    print("\n=== Testing Audio Size Limit ===")
    
    # Create a fake large audio file (over 25MB)
    large_audio = io.BytesIO(b"x" * (30 * 1024 * 1024))  # 30MB
    
    files = {
        'audio': ('large_audio.wav', large_audio, 'audio/wav')
    }
    data = {
        'source_lang': 'English'
    }
    
    try:
        response = requests.post(f"{BASE_URL}/transcribe", files=files, data=data)
        print(f"Status: {response.status_code}")
        if response.status_code == 413:
            print(f"✓ Correctly rejected large audio: {response.json()}")
        else:
            print(f"✗ Should have rejected large audio")
    except Exception as e:
        print(f"Error: {e}")

def test_valid_requests():
    """Test that valid-sized requests are accepted"""
    print("\n=== Testing Valid Size Requests ===")
    
    # Small JSON payload
    small_data = {
        "text": "Hello world",
        "source_lang": "English", 
        "target_lang": "Spanish"
    }
    
    try:
        response = requests.post(f"{BASE_URL}/translate", json=small_data)
        print(f"Small JSON - Status: {response.status_code}")
        if response.status_code != 413:
            print("✓ Small JSON accepted")
        else:
            print("✗ Small JSON should be accepted")
    except Exception as e:
        print(f"Error: {e}")
    
    # Small audio file
    small_audio = io.BytesIO(b"RIFF" + b"x" * 1000)  # 1KB fake WAV
    files = {
        'audio': ('small_audio.wav', small_audio, 'audio/wav')
    }
    data = {
        'source_lang': 'English'
    }
    
    try:
        response = requests.post(f"{BASE_URL}/transcribe", files=files, data=data)
        print(f"Small audio - Status: {response.status_code}")
        if response.status_code != 413:
            print("✓ Small audio accepted")
        else:
            print("✗ Small audio should be accepted")
    except Exception as e:
        print(f"Error: {e}")

def test_admin_endpoints():
    """Test admin endpoints for size limits"""
    print("\n=== Testing Admin Endpoints ===")
    
    headers = {'X-Admin-Token': os.environ.get('ADMIN_TOKEN', 'default-admin-token')}
    
    # Get current limits
    try:
        response = requests.get(f"{BASE_URL}/admin/size-limits", headers=headers)
        print(f"Get limits - Status: {response.status_code}")
        if response.status_code == 200:
            limits = response.json()
            print(f"✓ Current limits: {limits['limits_human']}")
        else:
            print(f"✗ Failed to get limits: {response.text}")
    except Exception as e:
        print(f"Error: {e}")
    
    # Update limits
    new_limits = {
        "max_audio_size": "30MB",
        "max_json_size": 2097152  # 2MB in bytes
    }
    
    try:
        response = requests.post(f"{BASE_URL}/admin/size-limits", 
                               json=new_limits, headers=headers)
        print(f"\nUpdate limits - Status: {response.status_code}")
        if response.status_code == 200:
            result = response.json()
            print(f"✓ Updated limits: {result['new_limits_human']}")
        else:
            print(f"✗ Failed to update limits: {response.text}")
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    print("Request Size Limit Tests")
    print("========================")
    print(f"Testing against: {BASE_URL}")
    print("\nMake sure the Flask app is running on port 5005")
    
    input("\nPress Enter to start tests...")
    
    test_valid_requests()
    test_json_size_limit()
    test_audio_size_limit()
    test_admin_endpoints()
    
    print("\n✅ All tests completed!")