test_api.py

"""
Test script for OmniParser API
"""
import requests
import json
import base64
from pathlib import Path

BASE_URL = "http://localhost:8000"


def health_check():
    """Check API health"""
    print("🏥 Health Check...")
    response = requests.get(f"{BASE_URL}/health")
    print(f"Status: {response.status_code}")
    print(json.dumps(response.json(), indent=2))
    print()


def parse_file(image_path: str):
    """Parse image file"""
    print(f"📸 Parsing file: {image_path}")
    
    if not Path(image_path).exists():
        print(f"❌ File not found: {image_path}")
        return
    
    with open(image_path, "rb") as f:
        files = {"file": f}
        response = requests.post(f"{BASE_URL}/parse", files=files)
    
    if response.status_code == 200:
        result = response.json()
        print(f"✅ Found {len(result['elements'])} UI elements")
        print(f"   Image size: {result['image_width']}x{result['image_height']}")
        print(f"   Processing time: {result['processing_time']:.2f}s")
        print("\n   Elements:")
        for elem in result['elements']:
            print(f"   - {elem['label']}: bbox={elem['bbox']}, confidence={elem['confidence']}")
    else:
        print(f"❌ Error: {response.status_code}")
        print(response.text)
    print()


def parse_base64(image_path: str):
    """Parse base64-encoded image"""
    print(f"📷 Parsing base64 image: {image_path}")
    
    if not Path(image_path).exists():
        print(f"❌ File not found: {image_path}")
        return
    
    # Read and encode image
    with open(image_path, "rb") as f:
        image_data = base64.b64encode(f.read()).decode('utf-8')
    
    payload = {
        "image_base64": image_data,
        "extract_text": True,
        "extract_icons": True
    }
    
    response = requests.post(f"{BASE_URL}/parse-base64", json=payload)
    
    if response.status_code == 200:
        result = response.json()
        print(f"✅ Found {len(result['elements'])} UI elements")
        print(f"   Processing time: {result['processing_time']:.2f}s")
    else:
        print(f"❌ Error: {response.status_code}")
        print(response.text)
    print()


if __name__ == "__main__":
    print("=" * 60)
    print("OmniParser API Test Suite")
    print("=" * 60)
    print()
    
    # Health check
    health_check()
    
    # Test with a sample image (if available)
    sample_images = [
        "screenshot.png",
        "test_image.png",
        "../screenshots/example.png"
    ]
    
    for img in sample_images:
        if Path(img).exists():
            parse_file(img)
            parse_base64(img)
            break
    else:
        print("⚠️  No test images found. Upload an image and try again.")
        print("   Expected: screenshot.png or test_image.png")
    
    print("=" * 60)
    print("✅ Test suite completed")
    print("=" * 60)