test_app.py

"""
Test script for BuildTheFuture application
Tests the core functionality without requiring API keys
"""

import os
import sys
from PIL import Image
import numpy as np

def create_test_image():
    """Create a simple test image for testing purposes"""
    # Create a simple construction site image
    width, height = 400, 300
    img = Image.new('RGB', (width, height), color='gray')
    
    # Add some basic construction elements
    pixels = np.array(img)
    
    # Add a foundation (dark gray rectangle)
    pixels[200:280, 50:350] = [64, 64, 64]
    
    # Add some steel beams (dark lines)
    pixels[180:190, 50:350] = [32, 32, 32]
    pixels[190:200, 50:350] = [32, 32, 32]
    
    # Add some incomplete walls (light gray)
    pixels[120:180, 100:150] = [128, 128, 128]
    pixels[120:180, 250:300] = [128, 128, 128]
    
    return Image.fromarray(pixels)

def test_image_processing():
    """Test basic image processing functionality"""
    print("🧪 Testing image processing...")
    
    try:
        # Create test image
        test_img = create_test_image()
        print("✅ Test image created successfully")
        
        # Test image properties
        print(f"   Image size: {test_img.size}")
        print(f"   Image mode: {test_img.mode}")
        
        # Save test image
        test_img.save("test_construction.jpg")
        print("✅ Test image saved as 'test_construction.jpg'")
        
        return True
    except Exception as e:
        print(f"❌ Image processing test failed: {e}")
        return False

def test_imports():
    """Test if all required modules can be imported"""
    print("🧪 Testing imports...")
    
    modules = [
        ("gradio", "gr"),
        ("PIL", "Image"),
        ("numpy", "np"),
        ("cv2", "cv2"),
        ("google.generativeai", "genai"),
        ("ultralytics", "YOLO"),
        ("requests", "requests"),
        ("elevenlabs", "generate"),
    ]
    
    failed_imports = []
    
    for module_name, alias in modules:
        try:
            if alias:
                exec(f"import {module_name} as {alias}")
            else:
                exec(f"import {module_name}")
            print(f"✅ {module_name}")
        except ImportError as e:
            print(f"❌ {module_name}: {e}")
            failed_imports.append(module_name)
    
    if failed_imports:
        print(f"\n⚠️  Failed imports: {', '.join(failed_imports)}")
        print("Run: pip install -r requirements.txt")
        return False
    
    print("✅ All imports successful")
    return True

def test_gradio_interface():
    """Test if Gradio interface can be created"""
    print("🧪 Testing Gradio interface...")
    
    try:
        import gradio as gr
        
        # Create a simple test interface
        def test_function(image):
            return image
        
        with gr.Blocks() as demo:
            gr.Markdown("# Test Interface")
            image_input = gr.Image()
            image_output = gr.Image()
            image_input.change(test_function, inputs=image_input, outputs=image_output)
        
        print("✅ Gradio interface created successfully")
        return True
    except Exception as e:
        print(f"❌ Gradio interface test failed: {e}")
        return False

def test_yolo_model():
    """Test YOLO model loading"""
    print("🧪 Testing YOLO model...")
    
    try:
        from ultralytics import YOLO
        
        # Try to load a small YOLO model
        model = YOLO('yolov11n.pt')
        print("✅ YOLO model loaded successfully")
        
        # Test with a simple image
        test_img = create_test_image()
        results = model(test_img)
        print("✅ YOLO inference test successful")
        
        return True
    except Exception as e:
        print(f"❌ YOLO model test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("🏗️  BuildTheFuture Test Suite")
    print("=" * 50)
    
    tests = [
        ("Import Test", test_imports),
        ("Image Processing Test", test_image_processing),
        ("Gradio Interface Test", test_gradio_interface),
        ("YOLO Model Test", test_yolo_model),
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n{test_name}:")
        print("-" * 30)
        if test_func():
            passed += 1
        print()
    
    print("=" * 50)
    print(f"Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! The application is ready to run.")
        print("\nNext steps:")
        print("1. Set up your API keys in .env file")
        print("2. Run: python app.py")
    else:
        print("⚠️  Some tests failed. Please fix the issues before running the application.")
    
    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)