scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

test(proxy): use real requests for IP allowlist checks

Browse source
This commit is contained in:
Devin AI 2026-11-03 00:23:24 +00:00
parent 67f4914d84
commit e619c19b01
1 changed files with 14 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -5,7 +5,7 @@
import litellm.proxy
import litellm.proxy.proxy_server
from typing import Dict, Final, List, Mapping, Optional
from typing import Final, List, Mapping, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -20,19 +20,12 @@ from litellm.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth,
get_api_key_from_custom_header,
)
from fastapi import WebSocket, HTTPException, status
from fastapi import Request, WebSocket, HTTPException, status
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles
from tests._master_key import MASTER_KEY
class Request:
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(
"allowed_ips, client_ip, expected_result",
[
@ -51,9 +44,11 @@ class Request:
def test_check_valid_ip(allowed_ips: Optional[List[str]], client_ip: Optional[str], expected_result: bool):
from litellm.proxy.auth.auth_utils import _check_valid_ip
request = Request(client_ip)
request: Final = Request(
{"type": "http", "client": (client_ip, 1234) if client_ip else None, "headers": []}
)
assert _check_valid_ip(allowed_ips, request)[0] == expected_result # type: ignore
assert _check_valid_ip(allowed_ips, request)[0] == expected_result
# test x-forwarder for is used when user has opted in
@ -79,9 +74,15 @@ def test_check_valid_ip_sent_with_x_forwarded_for(
):
from litellm.proxy.auth.auth_utils import _check_valid_ip
request = Request(client_ip, headers={"X-Forwarded-For": client_ip})
request: Final = Request(
{
"type": "http",
"client": ("192.0.2.1", 1234),
"headers": [(b"x-forwarded-for", client_ip.encode())] if client_ip else [],
}
)
assert _check_valid_ip(allowed_ips, request, use_x_forwarded_for=True)[0] == expected_result # type: ignore
assert _check_valid_ip(allowed_ips, request, use_x_forwarded_for=True)[0] == expected_result
@pytest.mark.asyncio