backend/test_api.py

#!/usr/bin/env python3
"""
Test script for Depth Flow Pro API
Tests all endpoints and measures performance
"""

import requests
import base64
import time
import sys
from pathlib import Path


class APITester:
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url
        self.test_results = []

    def print_header(self, text):
        """Print formatted header"""
        print("\n" + "=" * 60)
        print(f"  {text}")
        print("=" * 60)

    def print_result(self, name, success, time_ms=None, details=None):
        """Print test result"""
        status = "✓" if success else "✗"
        time_str = f" ({time_ms:.2f}ms)" if time_ms else ""
        print(f"{status} {name}{time_str}")

        if details:
            for key, value in details.items():
                print(f"    {key}: {value}")

        self.test_results.append({
            "name": name,
            "success": success,
            "time_ms": time_ms
        })

    def test_health(self):
        """Test health endpoint"""
        self.print_header("Testing Health Endpoints")

        try:
            # Root endpoint
            start = time.time()
            response = requests.get(f"{self.base_url}/")
            time_ms = (time.time() - start) * 1000

            data = response.json()
            self.print_result(
                "GET /",
                response.status_code == 200,
                time_ms,
                {"models_loaded": data.get("models_loaded", [])}
            )

            # Health endpoint
            start = time.time()
            response = requests.get(f"{self.base_url}/health")
            time_ms = (time.time() - start) * 1000

            data = response.json()
            self.print_result(
                "GET /health",
                response.status_code == 200,
                time_ms,
                {
                    "status": data.get("status"),
                    "gpu_enabled": data.get("gpu_enabled"),
                    "models": ", ".join(data.get("models", {}).keys())
                }
            )

            return True

        except Exception as e:
            print(f"✗ Health check failed: {e}")
            print("\nMake sure the server is running:")
            print("  cd backend")
            print("  uvicorn api.main:app --reload")
            return False

    def test_preview(self, image_path):
        """Test preview endpoint"""
        self.print_header("Testing Preview Endpoint")

        if not Path(image_path).exists():
            print(f"✗ Image not found: {image_path}")
            return False

        try:
            with open(image_path, 'rb') as f:
                start = time.time()

                response = requests.post(
                    f"{self.base_url}/api/v1/depth/preview",
                    files={'file': f}
                )

                time_ms = (time.time() - start) * 1000

            if response.status_code == 200:
                data = response.json()
                self.print_result(
                    "POST /api/v1/depth/preview",
                    True,
                    time_ms,
                    {
                        "model": data["metadata"]["model"],
                        "input_size": f"{data['metadata']['input_size'][0]}x{data['metadata']['input_size'][1]}",
                        "server_time": f"{data['processing_time_ms']:.2f}ms"
                    }
                )
                return True
            else:
                self.print_result(
                    "POST /api/v1/depth/preview",
                    False,
                    details={"error": response.text}
                )
                return False

        except Exception as e:
            print(f"✗ Preview test failed: {e}")
            return False

    def test_hq(self, image_path):
        """Test HQ endpoint"""
        self.print_header("Testing High Quality Endpoint")

        if not Path(image_path).exists():
            print(f"✗ Image not found: {image_path}")
            return False

        try:
            with open(image_path, 'rb') as f:
                start = time.time()

                response = requests.post(
                    f"{self.base_url}/api/v1/depth/hq",
                    files={'file': f}
                )

                time_ms = (time.time() - start) * 1000

            if response.status_code == 200:
                data = response.json()
                self.print_result(
                    "POST /api/v1/depth/hq",
                    True,
                    time_ms,
                    {
                        "model": data["metadata"]["model"],
                        "input_size": f"{data['metadata']['input_size'][0]}x{data['metadata']['input_size'][1]}",
                        "server_time": f"{data['processing_time_ms']:.2f}ms"
                    }
                )
                return True
            else:
                self.print_result(
                    "POST /api/v1/depth/hq",
                    False,
                    details={"error": response.text}
                )
                return False

        except Exception as e:
            print(f"✗ HQ test failed: {e}")
            return False

    def test_estimate(self, image_path):
        """Test custom estimate endpoint"""
        self.print_header("Testing Custom Estimate Endpoint")

        if not Path(image_path).exists():
            print(f"✗ Image not found: {image_path}")
            return False

        try:
            # Read and encode image
            with open(image_path, 'rb') as f:
                image_data = f.read()
                image_base64 = base64.b64encode(image_data).decode()

            # Test different configurations
            configs = [
                {"model": "small", "output_format": "grayscale", "colormap": "inferno"},
                {"model": "small", "output_format": "colormap", "colormap": "viridis"},
                {"model": "small", "output_format": "both", "colormap": "plasma"},
            ]

            for config in configs:
                start = time.time()

                response = requests.post(
                    f"{self.base_url}/api/v1/depth/estimate",
                    json={
                        "image": f"data:image/jpeg;base64,{image_base64}",
                        **config
                    }
                )

                time_ms = (time.time() - start) * 1000

                if response.status_code == 200:
                    data = response.json()
                    self.print_result(
                        f"Estimate ({config['output_format']})",
                        True,
                        time_ms,
                        {"colormap": config['colormap']}
                    )
                else:
                    self.print_result(
                        f"Estimate ({config['output_format']})",
                        False,
                        details={"error": response.text}
                    )

            return True

        except Exception as e:
            print(f"✗ Estimate test failed: {e}")
            return False

    def print_summary(self):
        """Print test summary"""
        self.print_header("Test Summary")

        total = len(self.test_results)
        passed = sum(1 for r in self.test_results if r["success"])
        failed = total - passed

        print(f"\nTotal Tests: {total}")
        print(f"Passed: {passed} ✓")
        print(f"Failed: {failed} ✗")

        if passed == total:
            print("\n🎉 All tests passed!")
        else:
            print(f"\n⚠️  {failed} test(s) failed")

        # Performance summary
        if any(r["time_ms"] for r in self.test_results):
            print("\nPerformance Summary:")
            for result in self.test_results:
                if result["time_ms"]:
                    print(f"  {result['name']}: {result['time_ms']:.2f}ms")


def main():
    """Run all tests"""
    print("""
    ╔═══════════════════════════════════════════════╗
    ║   Depth Flow Pro - API Test Suite            ║
    ╚═══════════════════════════════════════════════╝
    """)

    # Check for test image
    if len(sys.argv) < 2:
        print("Usage: python test_api.py <image_path>")
        print("\nExample:")
        print("  python test_api.py test_image.jpg")
        sys.exit(1)

    image_path = sys.argv[1]

    if not Path(image_path).exists():
        print(f"Error: Image not found: {image_path}")
        sys.exit(1)

    print(f"Test image: {image_path}")

    # Initialize tester
    tester = APITester()

    # Run tests
    if not tester.test_health():
        print("\n❌ Server not accessible. Aborting tests.")
        sys.exit(1)

    tester.test_preview(image_path)
    tester.test_hq(image_path)
    tester.test_estimate(image_path)

    # Print summary
    tester.print_summary()


if __name__ == "__main__":
    main()