TestGetCustomerIdFromStandardHeaders

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -1,9 +1,13 @@
 """
-Unit tests for auth_utils functions related to rate limiting.
+Unit tests for auth_utils functions related to rate limiting and customer ID extraction.
 """
 
+from unittest.mock import patch
+
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_utils import (
+    _get_customer_id_from_standard_headers,
+    get_end_user_id_from_request_body,
     get_key_model_rpm_limit,
     get_key_model_tpm_limit,
 )
@@ -129,3 +133,56 @@ class TestGetKeyModelTpmLimit:
         )
         result = get_key_model_tpm_limit(user_api_key_dict)
         assert result == {"gpt-4": 10000}
+
+
+class TestGetCustomerIdFromStandardHeaders:
+    """Tests for _get_customer_id_from_standard_headers helper function."""
+
+    def test_should_return_customer_id_from_x_litellm_customer_id_header(self):
+        """Should extract customer ID from x-litellm-customer-id header."""
+        headers = {"x-litellm-customer-id": "customer-123"}
+        result = _get_customer_id_from_standard_headers(request_headers=headers)
+        assert result == "customer-123"
+
+    def test_should_return_customer_id_from_x_litellm_end_user_id_header(self):
+        """Should extract customer ID from x-litellm-end-user-id header."""
+        headers = {"x-litellm-end-user-id": "end-user-456"}
+        result = _get_customer_id_from_standard_headers(request_headers=headers)
+        assert result == "end-user-456"
+
+    def test_should_return_none_when_headers_is_none(self):
+        """Should return None when headers is None."""
+        result = _get_customer_id_from_standard_headers(request_headers=None)
+        assert result is None
+
+    def test_should_return_none_when_no_standard_headers_present(self):
+        """Should return None when no standard customer ID headers are present."""
+        headers = {"x-other-header": "some-value"}
+        result = _get_customer_id_from_standard_headers(request_headers=headers)
+        assert result is None
+
+
+class TestGetEndUserIdFromRequestBodyWithStandardHeaders:
+    """Tests for get_end_user_id_from_request_body with standard customer ID headers."""
+
+    def test_should_prioritize_standard_header_over_body_user(self):
+        """Standard customer ID header should take precedence over body user field."""
+        headers = {"x-litellm-customer-id": "header-customer"}
+        request_body = {"user": "body-user"}
+
+        with patch("litellm.proxy.proxy_server.general_settings", {}):
+            result = get_end_user_id_from_request_body(
+                request_body=request_body, request_headers=headers
+            )
+        assert result == "header-customer"
+
+    def test_should_fall_back_to_body_when_no_standard_header(self):
+        """Should fall back to body user when no standard headers are present."""
+        headers = {"x-other-header": "value"}
+        request_body = {"user": "body-user"}
+
+        with patch("litellm.proxy.proxy_server.general_settings", {}):
+            result = get_end_user_id_from_request_body(
+                request_body=request_body, request_headers=headers
+            )
+        assert result == "body-user"