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

tests

Browse Source
This commit is contained in:
Luca Beurer-Kellner
2025-04-03 10:16:35 +02:00
parent 278b3a8bf4
commit b685830117
2 changed files with 94 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/common/request_context.py
+1 -2
View File
@@ -92,11 +92,10 @@ class RequestContext:
request_json=request_json,
dataset_name=dataset_name,
invariant_authorization=invariant_authorization,
guardrail_service_authorization=guardrail_service_authorization,
guardrail_authorization=guardrail_service_authorization,
guardrails=guardrails,
config=context_config,
_created_via_factory=True,
guardrail_authorization=guardrail_service_authorization,
)
def get_guardrailing_authorization(self) -> Optional[str]:
tests/unit_tests/common/test_authorization.py
+93 -11
View File
@@ -7,6 +7,9 @@ import random
import string
import pytest
from gateway.common.config_manager import GatewayConfig
from gateway.common.request_context import RequestContext
# Add root folder (parent) to sys.path
sys.path.append(
@@ -15,11 +18,19 @@ sys.path.append(
)
)
from gateway.common.authorization import extract_authorization_from_headers, INVARIANT_AUTHORIZATION_HEADER, API_KEYS_SEPARATOR
from gateway.common.authorization import (
INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER,
extract_authorization_from_headers,
INVARIANT_AUTHORIZATION_HEADER,
API_KEYS_SEPARATOR,
)
@pytest.mark.parametrize("push_to_explorer", [True, False])
@pytest.mark.parametrize("invariant_authorization", [True, False])
@pytest.mark.parametrize("invariant_authorization_appended_to_llm_provider_api_key", [True, False])
@pytest.mark.parametrize(
"invariant_authorization_appended_to_llm_provider_api_key", [True, False]
)
@pytest.mark.parametrize("use_fallback_header", [True, False])
def test_extract_authorization_from_headers(
push_to_explorer: bool,
@@ -29,19 +40,23 @@ def test_extract_authorization_from_headers(
):
"""Test the extract_authorization_from_headers function."""
llm_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
inv_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
llm_apikey = "".join(random.choices(string.ascii_letters + string.digits, k=10))
inv_apikey = "".join(random.choices(string.ascii_letters + string.digits, k=10))
dataset_name = "test-dataset" if push_to_explorer else None
headers: dict[str, str] = {}
headers: dict[str, str] = {}
llm_provider_api_key = "fallback-header" if use_fallback_header else "llm-provider-api-key"
llm_provider_api_key = (
"fallback-header" if use_fallback_header else "llm-provider-api-key"
)
headers[llm_provider_api_key] = llm_apikey
if invariant_authorization:
print("invariant_authorization - TRUE")
if invariant_authorization_appended_to_llm_provider_api_key:
headers[llm_provider_api_key] = f"{headers.get(llm_provider_api_key, '')}{API_KEYS_SEPARATOR}{inv_apikey}"
headers[llm_provider_api_key] = (
f"{headers.get(llm_provider_api_key, '')}{API_KEYS_SEPARATOR}{inv_apikey}"
)
else:
headers[INVARIANT_AUTHORIZATION_HEADER] = f"Bearer {inv_apikey}"
@@ -50,7 +65,7 @@ def test_extract_authorization_from_headers(
def __init__(self, headers):
self.headers = headers
request = MockRequest(headers)
request = MockRequest(headers)
# Call the function
try:
@@ -70,12 +85,79 @@ def test_extract_authorization_from_headers(
raise e
# Verify the results
if invariant_authorization:
if not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key:
if (
not push_to_explorer
and invariant_authorization_appended_to_llm_provider_api_key
):
assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[0] == llm_apikey
assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[1] == inv_apikey
else:
assert invariant_auth == ("Bearer " + inv_apikey)
else:
assert invariant_auth is None
if not(not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key):
assert llm_provider_api_key == llm_apikey
if not (
not push_to_explorer
and invariant_authorization_appended_to_llm_provider_api_key
):
assert llm_provider_api_key == llm_apikey
@pytest.mark.parametrize("use_guardrailing_api_key", [True, False])
def test_extract_guardrails_authorization_from_headers(use_guardrailing_api_key: bool):
headers: dict[str, str] = {}
inv_apikey = "".join(random.choices(string.ascii_letters + string.digits, k=10))
inv_guardrails_apikey = "".join(
random.choices(string.ascii_letters + string.digits, k=10)
)
llm_apikey = "".join(random.choices(string.ascii_letters + string.digits, k=10))
headers[INVARIANT_AUTHORIZATION_HEADER] = f"Bearer {inv_apikey}"
headers["Authorization"] = f"Bearer {llm_apikey}"
if use_guardrailing_api_key:
headers[INVARIANT_GUARDRAIL_SERVICE_AUTHORIZATION_HEADER] = (
f"Bearer {inv_guardrails_apikey}"
)
class MockRequest:
def __init__(self, headers):
self.headers = headers
dataset_name = "test-dataset"
request = MockRequest(headers)
try:
invariant_authorization, llm_provider_api_key = (
extract_authorization_from_headers(
request,
dataset_name=dataset_name,
llm_provider_api_key_header="Authorization",
)
)
context = RequestContext.create(
request_json={"input": "test"},
dataset_name=dataset_name,
invariant_authorization=invariant_authorization,
guardrails=None,
config=GatewayConfig(),
request=request,
)
except HTTPException as e:
# If an exception is raised, check if it is the expected one
if not invariant_authorization:
assert e.status_code == 400
assert e.detail == "Missing invariant api key"
return
else:
raise e
# Verify the results
assert invariant_authorization == ("Bearer " + inv_apikey)
assert llm_provider_api_key == llm_apikey
assert context.get_guardrailing_authorization() == (
"Bearer " + inv_guardrails_apikey
if use_guardrailing_api_key
else invariant_authorization
)