chopratejas--headroom/tests/test_google_multimodal_e2e.py

"""End-to-end test for Google Gemini multimodal content preservation.

This test uses the real Google Gemini API to verify that non-text content
(images, function calls) is preserved through the proxy's compression pipeline.

These tests require a GOOGLE_API_KEY environment variable and are skipped in CI.
Run manually with: GOOGLE_API_KEY=your_key python tests/test_google_multimodal_e2e.py
"""

import asyncio
import os

import httpx
import pytest

# 10x10 red pixel PNG for testing (valid image generated by PIL)
TINY_RED_PNG = "iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAEklEQVR4nGP8z4APMOGVHbHSAEEsAROxCnMTAAAAAElFTkSuQmCC"


@pytest.fixture
def api_key():
    """Get API key from environment, skip if not available."""
    key = os.environ.get("GOOGLE_API_KEY")
    if not key:
        pytest.skip("GOOGLE_API_KEY not set - skipping E2E tests")
    return key


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_text_only_request(api_key):
    """Test that pure text requests work normally."""
    print("\n=== Test 1: Pure Text Request ===")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"

    payload = {
        "contents": [
            {"role": "user", "parts": [{"text": "What is 2 + 2? Reply with just the number."}]}
        ]
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload, timeout=30)

    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        text = (
            data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
        )
        print(f"Response: {text[:100]}")
        print("✅ Text-only request works")
        return True
    else:
        print(f"Error: {response.text[:200]}")
        return False


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_image_request(api_key):
    """Test that image content is preserved and processed."""
    print("\n=== Test 2: Image Request (inlineData) ===")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"

    # Request with inline image
    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [
                    {"text": "What color is this tiny image? Reply with just the color name."},
                    {"inlineData": {"mimeType": "image/png", "data": TINY_RED_PNG}},
                ],
            }
        ]
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload, timeout=30)

    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        text = (
            data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
        )
        print(f"Response: {text[:100]}")
        print("✅ Image request works - model processed the image")
        return True
    else:
        print(f"Error: {response.text[:200]}")
        return False


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_function_calling(api_key):
    """Test that function calling works (functionCall in response)."""
    print("\n=== Test 3: Function Calling ===")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"

    # Request with function declaration
    payload = {
        "contents": [{"role": "user", "parts": [{"text": "What's the weather in New York?"}]}],
        "tools": [
            {
                "functionDeclarations": [
                    {
                        "name": "get_weather",
                        "description": "Get the current weather for a location",
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "location": {"type": "string", "description": "The city name"}
                            },
                            "required": ["location"],
                        },
                    }
                ]
            }
        ],
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload, timeout=30)

    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])

        # Check if model made a function call
        has_function_call = any("functionCall" in part for part in parts)
        if has_function_call:
            func_call = next(p["functionCall"] for p in parts if "functionCall" in p)
            print(f"Function called: {func_call.get('name')} with args: {func_call.get('args')}")
            print("✅ Function calling works")
            return True
        else:
            # Model might have answered directly
            text = parts[0].get("text", "") if parts else ""
            print(f"Model responded with text instead: {text[:100]}")
            print("⚠️ Model didn't use function call (acceptable)")
            return True
    else:
        print(f"Error: {response.text[:200]}")
        return False


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_function_response_flow(api_key):
    """Test complete function call + response flow."""
    print("\n=== Test 4: Function Call + Response Flow ===")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"

    # Multi-turn with function response
    payload = {
        "contents": [
            {"role": "user", "parts": [{"text": "What's the weather in Tokyo?"}]},
            {
                "role": "model",
                "parts": [{"functionCall": {"name": "get_weather", "args": {"location": "Tokyo"}}}],
            },
            {
                "role": "user",
                "parts": [
                    {
                        "functionResponse": {
                            "name": "get_weather",
                            "response": {"temperature": 22, "condition": "sunny", "humidity": 45},
                        }
                    }
                ],
            },
        ],
        "tools": [
            {
                "functionDeclarations": [
                    {
                        "name": "get_weather",
                        "description": "Get the current weather for a location",
                        "parameters": {
                            "type": "object",
                            "properties": {"location": {"type": "string"}},
                        },
                    }
                ]
            }
        ],
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload, timeout=30)

    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        text = (
            data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
        )
        print(f"Response: {text[:150]}")
        print("✅ Function response flow works - model used the function result")
        return True
    else:
        print(f"Error: {response.text[:300]}")
        return False


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_mixed_conversation(api_key):
    """Test a conversation mixing text and images."""
    print("\n=== Test 5: Mixed Conversation (Text + Image) ===")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"

    payload = {
        "contents": [
            {"role": "user", "parts": [{"text": "I'll show you an image and ask about it."}]},
            {
                "role": "model",
                "parts": [{"text": "Sure, please share the image and I'll help you with it."}],
            },
            {
                "role": "user",
                "parts": [
                    {"text": "Here it is. What color do you see?"},
                    {"inlineData": {"mimeType": "image/png", "data": TINY_RED_PNG}},
                ],
            },
        ]
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload, timeout=30)

    print(f"Status: {response.status_code}")
    if response.status_code == 200:
        data = response.json()
        text = (
            data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
        )
        print(f"Response: {text[:150]}")
        print("✅ Mixed conversation works - model saw and processed the image")
        return True
    else:
        print(f"Error: {response.text[:200]}")
        return False


@pytest.mark.skipif(
    not os.environ.get("GOOGLE_API_KEY"),
    reason="GOOGLE_API_KEY not set - E2E tests require real API access",
)
@pytest.mark.asyncio
async def test_through_proxy(api_key, proxy_url: str = "http://localhost:8080"):
    """Test multimodal requests through the Headroom proxy."""
    print(f"\n=== Test 6: Through Headroom Proxy ({proxy_url}) ===")

    # The proxy expects requests at /v1beta/models/{model}:generateContent
    url = f"{proxy_url}/v1beta/models/gemini-2.0-flash:generateContent"

    payload = {
        "contents": [
            {
                "role": "user",
                "parts": [
                    {"text": "Describe this image in one word."},
                    {"inlineData": {"mimeType": "image/png", "data": TINY_RED_PNG}},
                ],
            }
        ]
    }

    headers = {"x-goog-api-key": api_key, "Content-Type": "application/json"}

    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload, headers=headers, timeout=30)

        print(f"Status: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            text = (
                data.get("candidates", [{}])[0]
                .get("content", {})
                .get("parts", [{}])[0]
                .get("text", "")
            )
            print(f"Response: {text[:150]}")
            print("✅ Proxy preserved the image and forwarded correctly!")
            return True
        else:
            print(f"Error: {response.text[:300]}")
            return False
    except httpx.ConnectError:
        print("⚠️ Proxy not running - skipping proxy test")
        print("   To test through proxy, start it with: uv run headroom-proxy")
        return None


async def main():
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        print("ERROR: GOOGLE_API_KEY environment variable not set")
        print("Usage: GOOGLE_API_KEY=your_key python tests/test_google_multimodal_e2e.py")
        return False

    print("=" * 60)
    print("Google Gemini Multimodal E2E Tests")
    print("=" * 60)
    print(f"Using API key: {api_key[:10]}...")

    results = []

    # Test 1: Pure text
    results.append(("Text Only", await test_text_only_request(api_key)))

    # Test 2: Image
    results.append(("Image (inlineData)", await test_image_request(api_key)))

    # Test 3: Function calling
    results.append(("Function Calling", await test_function_calling(api_key)))

    # Test 4: Function response
    results.append(("Function Response Flow", await test_function_response_flow(api_key)))

    # Test 5: Mixed conversation
    results.append(("Mixed Conversation", await test_mixed_conversation(api_key)))

    # Test 6: Through proxy (if running)
    proxy_result = await test_through_proxy(api_key)
    if proxy_result is not None:
        results.append(("Through Proxy", proxy_result))

    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)

    passed = sum(1 for _, r in results if r)
    total = len(results)

    for name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"  {name}: {status}")

    print(f"\nTotal: {passed}/{total} passed")

    return passed == total


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