Refactor tests.

2026-02-12 22:42:45 +00:00 · 2025-04-01 15:15:32 +02:00
parent 750c83d3f8
commit eced3755b2
13 changed files with 121 additions and 211 deletions
--- a/run.sh
+++ b/run.sh
@@ -98,7 +98,7 @@ integration_tests() {
    exit 1
  fi

-  TEST_GUARDRAILS_FILE_PATH="tests/integration/resources/guardrails/find_capital_guardrails.py"
+  TEST_GUARDRAILS_FILE_PATH="tests/integration/resources/guardrails/integration_test_guardrails_via_file.py"
  if [[ -n "$TEST_GUARDRAILS_FILE_PATH" ]]; then
    if [[ -f "$TEST_GUARDRAILS_FILE_PATH" ]]; then
      TEST_GUARDRAILS_FILE_PATH=$(realpath "$TEST_GUARDRAILS_FILE_PATH")
--- a/tests/integration/anthropic/test_anthropic_with_tool_call.py
+++ b/tests/integration/anthropic/test_anthropic_with_tool_call.py
@@ -12,10 +12,11 @@ from typing import Dict, List
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_anthropic_client
+
 import anthropic
 import pytest
 import requests
-from httpx import Client

 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -26,14 +27,8 @@ class WeatherAgent:

    def __init__(self, gateway_url, push_to_explorer):
        self.dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-        invariant_api_key = os.environ.get("INVARIANT_API_KEY")
-        self.client = anthropic.Anthropic(
-            http_client=Client(
-                headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-            ),
-            base_url=f"{gateway_url}/api/v1/gateway/{self.dataset_name}/anthropic"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/anthropic",
+        self.client = get_anthropic_client(
+            gateway_url, push_to_explorer, self.dataset_name
        )
        self.get_weather_function = {
            "name": "get_weather",
--- a/tests/integration/anthropic/test_anthropic_without_tool_call.py
+++ b/tests/integration/anthropic/test_anthropic_without_tool_call.py
@@ -8,10 +8,10 @@ import uuid
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

-import anthropic
+from utils import get_anthropic_client
+
 import pytest
 import requests
-from httpx import Client

 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -26,15 +26,10 @@ async def test_response_without_tool_call(
 ):
    """Test the Anthropic gateway without tool calling."""
    dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-    invariant_api_key = os.environ.get("INVARIANT_API_KEY")
-
-    client = anthropic.Anthropic(
-        http_client=Client(
-            headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    cities = ["zurich", "new york", "london"]
@@ -91,16 +86,7 @@ async def test_streaming_response_without_tool_call(
 ):
    """Test the Anthropic gateway without tool calling."""
    dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-    invariant_api_key = os.environ.get("INVARIANT_API_KEY")
-
-    client = anthropic.Anthropic(
-        http_client=Client(
-            headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
-    )
+    client = get_anthropic_client(gateway_url, push_to_explorer, dataset_name)

    cities = ["zurich", "new york", "london"]
    queries = [
--- a/tests/integration/docker-compose.test.yml
+++ b/tests/integration/docker-compose.test.yml
@@ -60,6 +60,7 @@ services:
  app-api:
    container_name: invariant-gateway-test-explorer-app-api
    image: ghcr.io/invariantlabs-ai/explorer/app-api:latest
+    pull_policy: always
    platform: linux/amd64
    depends_on:
      database:
--- a/tests/integration/gemini/test_generate_content_with_tool_calls.py
+++ b/tests/integration/gemini/test_generate_content_with_tool_calls.py
@@ -8,9 +8,10 @@ import uuid
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_gemini_client
+
 import pytest
 import requests
-from google import genai
 from google.genai import types

 # Pytest plugins
@@ -143,18 +144,7 @@ async def test_generate_content_with_tool_call(
    without streaming.
    """
    dataset_name = f"test-dataset-gemini-{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": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        },
-    )
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)

    request = {
        "model": "gemini-2.0-flash",
--- a/tests/integration/gemini/test_generate_content_without_tool_calls.py
+++ b/tests/integration/gemini/test_generate_content_without_tool_calls.py
@@ -10,6 +10,8 @@ from unittest.mock import patch
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_gemini_client
+
 import pytest
 import PIL.Image
 import requests
@@ -29,17 +31,8 @@ async def test_generate_content(
 ):
    """Test the generate content gateway calls without tool calling."""
    dataset_name = f"test-dataset-gemini-{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": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        },
-    )
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)
+
    request = {
        "model": "gemini-2.0-flash",
        "contents": "What is the capital of France?",
@@ -115,18 +108,8 @@ async def test_generate_content_with_image(
 ):
    """Test that generate content gateway calls work with image."""
    dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)

-    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": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        },
-    )

    image_path = Path(__file__).parent.parent / "resources" / "images" / "two-cats.png"
    image = PIL.Image.open(image_path)
--- a/tests/integration/guardrails/test_guardrails_anthropic.py
+++ b/tests/integration/guardrails/test_guardrails_anthropic.py
@@ -8,6 +8,8 @@ import time
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_anthropic_client
+
 import pytest
 import requests
 from httpx import Client
@@ -32,16 +34,10 @@ async def test_message_content_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    request = {
@@ -161,16 +157,10 @@ async def test_tool_call_guardrail_from_file(
    }

    dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    if not do_stream:
@@ -255,16 +245,10 @@ async def test_input_from_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    request = {
--- a/tests/integration/guardrails/test_guardrails_gemini.py
+++ b/tests/integration/guardrails/test_guardrails_gemini.py
@@ -8,9 +8,10 @@ import time
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_gemini_client
+
 import pytest
 import requests
-from httpx import Client
 from google import genai

 # Pytest plugins
@@ -30,17 +31,10 @@ async def test_message_content_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    request = {
@@ -141,17 +135,10 @@ async def test_tool_call_guardrail_from_file(
    )

    dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    request = {
@@ -244,17 +231,10 @@ async def test_input_from_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
    )

    request = {
--- a/tests/integration/guardrails/test_guardrails_open_ai.py
+++ b/tests/integration/guardrails/test_guardrails_open_ai.py
@@ -8,6 +8,8 @@ import time
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_open_ai_client
+
 import pytest
 import requests
 from httpx import Client
@@ -30,17 +32,7 @@ async def test_message_content_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    request = {
        "model": "gpt-4o",
@@ -161,17 +153,7 @@ async def test_tool_call_guardrail_from_file(
    }

    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    if not do_stream:
        with pytest.raises(BadRequestError) as exc_info:
@@ -259,17 +241,7 @@ async def test_input_from_guardrail_from_file(
        pytest.fail("No INVARIANT_API_KEY set, failing")

    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    request = {
        "model": "gpt-4o",
--- a/tests/integration/open_ai/test_chat_with_tool_call.py
+++ b/tests/integration/open_ai/test_chat_with_tool_call.py
@@ -9,10 +9,10 @@ import uuid
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_open_ai_client
+
 import pytest
 import requests
-from httpx import Client
-from openai import OpenAI

 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -28,17 +28,7 @@ async def test_chat_completion_with_tool_call_without_streaming(
    without streaming.
    """
    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    chat_response = client.chat.completions.create(
        model="gpt-4o",
@@ -146,17 +136,7 @@ async def test_chat_completion_with_tool_call_with_streaming(
    while streaming.
    """
    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    chat_response = client.chat.completions.create(
        model="gpt-4o",
--- a/tests/integration/open_ai/test_chat_without_tool_calls.py
+++ b/tests/integration/open_ai/test_chat_without_tool_calls.py
@@ -11,6 +11,8 @@ from unittest.mock import patch
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

+from utils import get_open_ai_client
+
 import pytest
 import requests
 from httpx import Client
@@ -30,17 +32,7 @@ async def test_chat_completion(
 ):
    """Test the chat completions gateway calls without tool calling."""
    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

    chat_response = client.chat.completions.create(
        model="gpt-4o",
@@ -103,17 +95,8 @@ async def test_chat_completion_with_image(
 ):
    """Test the chat completions gateway works with image."""
    dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)

-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
    image_path = Path(__file__).parent.parent / "resources" / "images" / "two-cats.png"
    with image_path.open("rb") as image_file:
        base64_image = base64.b64encode(image_file.read()).decode("utf-8")
--- a/tests/integration/resources/guardrails/integration_test_guardrails_via_file.py
+++ b/tests/integration/resources/guardrails/integration_test_guardrails_via_file.py
@@ -1,15 +1,15 @@
-# For Anthropic and OpenAI tests:
+# For Anthropic and OpenAI tests
 raise "Madrid detected in the response" if:
    (msg: Message)
   "Madrid" in msg.content

-# For Gemini tests:
+# For Gemini tests
 raise "Dublin detected in the response" if:
    (msg: Message)
    (chunk: str) in text(msg.content)
    "Dublin" in chunk

-# For all provider tests:
+# For all provider tests
 raise "get_capital is called with Germany as argument" if:
    (call: ToolCall)
    call is tool:get_capital
--- a/tests/integration/utils.py
+++ b/tests/integration/utils.py
@@ -0,0 +1,56 @@
+"""Common utilities for integration tests."""
+
+import os
+from httpx import Client
+from openai import OpenAI
+from google import genai
+from anthropic import Anthropic
+
+
+def get_open_ai_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> OpenAI:
+    """Create an OpenAI client for integration tests."""
+    return OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/openai",
+    )
+
+
+def get_anthropic_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> Anthropic:
+    """Create an Anthropic client for integration tests."""
+    return Anthropic(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/anthropic",
+    )
+
+
+def get_gemini_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> genai.Client:
+    """Create a Gemini client for integration tests."""
+    return 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": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        },
+    )