test_error_handling.py

"""
Test script to verify error handling in the object detection functionality.
"""

from PIL import Image
import pandas as pd
import app
from app import detect_objects, initialize_model

def test_with_none_model():
    """Test detection when model is None."""
    print("Testing with None model...")
    
    # Temporarily set model to None
    original_model = app.model
    app.model = None
    
    # Create test image
    test_img = Image.new('RGB', (100, 100), 'white')
    
    # Run detection
    annotated_img, detection_df = detect_objects(test_img)
    
    # Verify results
    print(f"Returned image type: {type(annotated_img)}")
    print(f"Returned DataFrame shape: {detection_df.shape}")
    print(f"DataFrame columns: {list(detection_df.columns)}")
    print(f"DataFrame empty: {detection_df.empty}")
    
    # Restore original model
    app.model = original_model
    print("Model restored\n")

def test_with_corrupted_image():
    """Test detection with edge case images."""
    print("Testing with very small image...")
    
    # Create very small image
    tiny_img = Image.new('RGB', (1, 1), 'black')
    
    # Run detection
    annotated_img, detection_df = detect_objects(tiny_img)
    
    print(f"Small image - Detections: {len(detection_df)}")
    expected_columns = ['class', 'confidence', 'x1', 'y1', 'x2', 'y2']
    print(f"Small image - DataFrame columns: {list(detection_df.columns)}")
    print(f"Small image - DataFrame columns correct: {list(detection_df.columns) == expected_columns}")
    
    print("Testing with large image...")
    
    # Create large image
    large_img = Image.new('RGB', (2000, 2000), 'blue')
    
    # Run detection
    annotated_img, detection_df = detect_objects(large_img)
    
    print(f"Large image - Detections: {len(detection_df)}")
    print(f"Large image - DataFrame columns: {list(detection_df.columns)}")
    print(f"Large image - DataFrame columns correct: {list(detection_df.columns) == expected_columns}")

def main():
    print("Testing error handling scenarios...\n")
    
    # Initialize model first
    if not initialize_model():
        print("Failed to initialize model for testing")
        return
    
    test_with_none_model()
    test_with_corrupted_image()
    
    print("Error handling tests completed!")

if __name__ == "__main__":
    main()