mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Add E2E tests for Google multimodal content preservation
- Tests for text-only, image, function calling, function response, mixed conversation - Tests skip automatically when GOOGLE_API_KEY not set (CI-safe) - Can run standalone: GOOGLE_API_KEY=key python tests/test_google_multimodal_e2e.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
chopratejas
parent
dd4c65e565
commit
62ee1dc9e3
1 changed files with 375 additions and 0 deletions
375
tests/test_google_multimodal_e2e.py
Normal file
375
tests/test_google_multimodal_e2e.py
Normal file
|
|
@ -0,0 +1,375 @@
|
|||
"""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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue