diff --git a/gateway/common/request_context.py b/gateway/common/request_context.py index 42ffc45..618cf4d 100644 --- a/gateway/common/request_context.py +++ b/gateway/common/request_context.py @@ -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]: diff --git a/tests/unit_tests/common/test_authorization.py b/tests/unit_tests/common/test_authorization.py index 3742f35..60512aa 100644 --- a/tests/unit_tests/common/test_authorization.py +++ b/tests/unit_tests/common/test_authorization.py @@ -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 \ No newline at end of file + 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 + )