mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
274bf3e48d
* fix get_response_headers * unit testing for get headers * unit testing for anthropic / azure openai headers * increase test coverage for test_completion_response_ratelimit_headers * fix test rate limit headers
96 lines
3 KiB
Python
96 lines
3 KiB
Python
import sys
|
|
import os
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../")
|
|
) # Adds the parent directory to the system path
|
|
|
|
import pytest
|
|
from litellm.llms.AzureOpenAI.common_utils import process_azure_headers
|
|
from httpx import Headers
|
|
|
|
|
|
def test_process_azure_headers_empty():
|
|
result = process_azure_headers({})
|
|
assert result == {}, "Expected empty dictionary for no input"
|
|
|
|
|
|
def test_process_azure_headers_with_all_headers():
|
|
input_headers = Headers(
|
|
{
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-requests": "90",
|
|
"x-ratelimit-limit-tokens": "10000",
|
|
"x-ratelimit-remaining-tokens": "9000",
|
|
"other-header": "value",
|
|
}
|
|
)
|
|
|
|
expected_output = {
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-requests": "90",
|
|
"x-ratelimit-limit-tokens": "10000",
|
|
"x-ratelimit-remaining-tokens": "9000",
|
|
"llm_provider-x-ratelimit-limit-requests": "100",
|
|
"llm_provider-x-ratelimit-remaining-requests": "90",
|
|
"llm_provider-x-ratelimit-limit-tokens": "10000",
|
|
"llm_provider-x-ratelimit-remaining-tokens": "9000",
|
|
"llm_provider-other-header": "value",
|
|
}
|
|
|
|
result = process_azure_headers(input_headers)
|
|
assert result == expected_output, "Unexpected output for all Azure headers"
|
|
|
|
|
|
def test_process_azure_headers_with_partial_headers():
|
|
input_headers = Headers(
|
|
{
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-tokens": "9000",
|
|
"other-header": "value",
|
|
}
|
|
)
|
|
|
|
expected_output = {
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-tokens": "9000",
|
|
"llm_provider-x-ratelimit-limit-requests": "100",
|
|
"llm_provider-x-ratelimit-remaining-tokens": "9000",
|
|
"llm_provider-other-header": "value",
|
|
}
|
|
|
|
result = process_azure_headers(input_headers)
|
|
assert result == expected_output, "Unexpected output for partial Azure headers"
|
|
|
|
|
|
def test_process_azure_headers_with_no_matching_headers():
|
|
input_headers = Headers(
|
|
{"unrelated-header-1": "value1", "unrelated-header-2": "value2"}
|
|
)
|
|
|
|
expected_output = {
|
|
"llm_provider-unrelated-header-1": "value1",
|
|
"llm_provider-unrelated-header-2": "value2",
|
|
}
|
|
|
|
result = process_azure_headers(input_headers)
|
|
assert result == expected_output, "Unexpected output for non-matching headers"
|
|
|
|
|
|
def test_process_azure_headers_with_dict_input():
|
|
input_headers = {
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-requests": "90",
|
|
"other-header": "value",
|
|
}
|
|
|
|
expected_output = {
|
|
"x-ratelimit-limit-requests": "100",
|
|
"x-ratelimit-remaining-requests": "90",
|
|
"llm_provider-x-ratelimit-limit-requests": "100",
|
|
"llm_provider-x-ratelimit-remaining-requests": "90",
|
|
"llm_provider-other-header": "value",
|
|
}
|
|
|
|
result = process_azure_headers(input_headers)
|
|
assert result == expected_output, "Unexpected output for dict input"
|