CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-07 11:27:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add tests for some non tool_call cases for Gemini integration.

Browse Source
This commit is contained in:
Hemang
2025-03-10 13:18:33 +01:00
committed by Hemang Sarkar
parent 5a6dd5c813
commit fdc9da16c0
3 changed files with 238 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
run.sh
+14 -2
View File
@@ -50,7 +50,7 @@ down() {
tests() {
echo "Setting up test environment..."
echo "Setting up test environment to run integration tests..."
# Ensure test network exists
docker network inspect invariant-gateway-web-test >/dev/null 2>&1 || \
@@ -84,7 +84,18 @@ tests() {
sleep 2
done
echo "Running tests..."
while true; do
echo "Attempting to create test user in invariant-gateway-test-explorer-app-api"
RESPONSE=$(curl -ks -X POST http://127.0.0.1/api/v1/user/signup)
if echo "$RESPONSE" | jq -e '.success == true' >/dev/null 2>&1; then
echo "Created test user in invariant-gateway-test-explorer-app-api"
break
fi
echo "$RESPONSE"
sleep 2
done
echo "Running integration tests..."
# Make call to signup endpoint
curl -k -X POST http://127.0.0.1/api/v1/user/signup
@@ -96,6 +107,7 @@ tests() {
--network invariant-gateway-web-test \
-e OPENAI_API_KEY="$OPENAI_API_KEY" \
-e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\
-e GEMINI_API_KEY="$GEMINI_API_KEY" \
--env-file ./tests/.env.test \
invariant-gateway-tests $@
}
tests/gemini/test_generate_content_without_tool_calls.py
+222
View File
@@ -0,0 +1,222 @@
"""Test the generate content gateway calls without tool calling."""
import base64
import os
import sys
import uuid
from pathlib import Path
from unittest.mock import patch
import pytest
from httpx import Client
from google import genai
import PIL.Image
# add tests folder (parent) to sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from util import * # needed for pytest fixtures
pytest_plugins = ("pytest_asyncio",)
@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
@pytest.mark.parametrize(
"do_stream, push_to_explorer",
[(True, True), (True, False), (False, True), (False, False)],
)
async def test_generate_content(
context, explorer_api_url, gateway_url, do_stream, push_to_explorer
):
"""Test the generate content gateway calls without tool calling."""
dataset_name = "test-dataset-gemini-" + str(uuid.uuid4())
client = genai.Client(
api_key=os.getenv("GEMINI_API_KEY"),
http_options={
"base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
if push_to_explorer
else f"{gateway_url}/api/v1/gateway/gemini",
"headers": {
"invariant-authorization": "Bearer <some-key>"
}, # This key is not used for local tests
},
)
request = {
"model": "gemini-2.0-flash",
"contents": "What is the capital of France?",
"config": {
"maxOutputTokens": 100,
"system_instruction": "This is the system instruction.",
},
}
if not do_stream:
chat_response = client.models.generate_content(**request)
else:
chat_response = client.models.generate_content_stream(**request)
# Verify the chat response
if not do_stream:
assert "PARIS" in chat_response.candidates[0].content.parts[0].text.upper()
expected_assistant_message = chat_response.candidates[0].content.parts[0].text
else:
full_response = ""
for chunk in chat_response:
if (
chunk.candidates
and chunk.candidates[0].content
and chunk.candidates[0].content.parts
):
for text_part in chunk.candidates[0].content.parts:
full_response += text_part.text
assert "PARIS" in full_response.upper()
expected_assistant_message = full_response
if push_to_explorer:
# Fetch the trace ids for the dataset
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == 1
trace_id = traces[0]["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
# Verify the trace messages
assert trace["messages"] == [
{
"role": "system",
"content": "This is the system instruction.",
},
{
"role": "user",
"content": [{"text": "What is the capital of France?", "type": "text"}],
},
{
"role": "assistant",
"content": expected_assistant_message,
},
]
@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
@pytest.mark.parametrize("push_to_explorer", [True, False])
async def test_generate_content_with_image(
context, explorer_api_url, gateway_url, push_to_explorer
):
"""Test that generate content gateway calls work with image."""
dataset_name = "test-dataset-gemini-" + str(uuid.uuid4())
client = genai.Client(
api_key=os.getenv("GEMINI_API_KEY"),
http_options={
"base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
if push_to_explorer
else f"{gateway_url}/api/v1/gateway/gemini",
"headers": {
"invariant-authorization": "Bearer <some-key>"
}, # This key is not used for local tests
},
)
image_path = Path(__file__).parent.parent / "images" / "two-cats.png"
image = PIL.Image.open(image_path)
chat_response = client.models.generate_content(
model="gemini-2.0-flash",
contents=["How many cats are there in this image?", image],
config={"maxOutputTokens": 100},
)
assert "TWO" in chat_response.candidates[0].content.parts[0].text.upper()
if push_to_explorer:
# Fetch the trace ids for the dataset
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == 1
trace_id = traces[0]["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
# Verify the trace messages
assert len(trace["messages"]) == 2
assert trace["messages"][0]["role"] == "user"
assert trace["messages"][0]["content"][0] == {
"type": "text",
"text": "How many cats are there in this image?",
}
assert trace["messages"][0]["content"][1]["type"] == "image_url"
assert trace["messages"][1] == {
"role": "assistant",
"content": chat_response.candidates[0].content.parts[0].text,
}
@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
async def test_generate_content_with_invariant_key_in_gemini_key_header(
context, explorer_api_url, gateway_url
):
"""Test the generate content gateway calls with the Invariant API Key in the Gemini Key header."""
dataset_name = "test-dataset-gemini-" + str(uuid.uuid4())
gemini_api_key = os.getenv("GEMINI_API_KEY")
with patch.dict(
os.environ,
{"GEMINI_API_KEY": gemini_api_key + ";invariant-auth=<not needed for test>"},
):
client = genai.Client(
api_key=os.getenv("GEMINI_API_KEY"),
http_options={
"base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
},
)
chat_response = client.models.generate_content(
model="gemini-2.0-flash",
contents="What is the capital of Spain?",
config={
"maxOutputTokens": 100,
},
)
# Verify the chat response
assert "MADRID" in chat_response.candidates[0].content.parts[0].text.upper()
expected_assistant_message = chat_response.candidates[0].content.parts[0].text
# Fetch the trace ids for the dataset
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == 1
trace_id = traces[0]["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
# Verify the trace messages
assert trace["messages"] == [
{
"role": "user",
"content": [{"text": "What is the capital of Spain?", "type": "text"}],
},
{
"role": "assistant",
"content": expected_assistant_message,
},
]
tests/requirements.txt
+2
View File
@@ -1,5 +1,7 @@
anthropic
google-genai
openai
pillow
pytest
pytest-asyncio
pytest-playwright