phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add test for test_check_valid_ip_sent_with_x_forwarded_for

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-13 15:54:53 -07:00
parent b94c982ec9
commit d0e6ca659f
2 changed files with 32 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1221,7 +1221,7 @@ def _check_valid_ip(
# if general_settings.get("use_x_forwarded_for") is True then use x-forwarded-for
client_ip = None
if use_x_forwarded_for is True:
if use_x_forwarded_for is True and "x-forwarded-for" in request.headers:
client_ip = request.headers["x-forwarded-for"]
elif request.client is not None:
client_ip = request.client.host

33
litellm/tests/test_user_api_key_auth.py
View file

@ -7,7 +7,7 @@ import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from typing import List, Optional
from typing import Dict, List, Optional
from unittest.mock import MagicMock
import pytest
@ -16,9 +16,10 @@ import litellm
class Request:
def __init__(self, client_ip: Optional[str] = None):
def __init__(self, client_ip: Optional[str] = None, headers: Optional[dict] = None):
self.client = MagicMock()
self.client.host = client_ip
self.headers: Dict[str, str] = {}
@pytest.mark.parametrize(
@ -46,6 +47,34 @@ def test_check_valid_ip(
assert _check_valid_ip(allowed_ips, request) == expected_result # type: ignore
# test x-forwarder for is used when user has opted in
@pytest.mark.parametrize(
"allowed_ips, client_ip, expected_result",
[
(None, "127.0.0.1", True), # No IP restrictions, should be allowed
(["127.0.0.1"], "127.0.0.1", True), # IP in allowed list
(["192.168.1.1"], "127.0.0.1", False), # IP not in allowed list
([], "127.0.0.1", False), # Empty allowed list, no IP should be allowed
(["192.168.1.1", "10.0.0.1"], "10.0.0.1", True), # IP in allowed list
(
["192.168.1.1"],
None,
False,
), # Request with no client IP should not be allowed
],
)
def test_check_valid_ip_sent_with_x_forwarded_for(
allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool
):
from litellm.proxy.auth.user_api_key_auth import _check_valid_ip
request = Request(client_ip, headers={"X-Forwarded-For": client_ip})
assert _check_valid_ip(allowed_ips, request, use_x_forwarded_for=True) == expected_result # type: ignore
@pytest.mark.asyncio
async def test_check_blocked_team():
"""