From 844d6828c86630452b0b85137c835a87e54ed7ee Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 16 Mar 2024 13:53:10 -0700
Subject: [PATCH] fix(proxy_server.py): rename beta end user blocking endpoints

separating end user management from user management
---
 litellm/proxy/proxy_server.py           |  5 +-
 litellm/tests/test_blocked_user_list.py | 91 ++++++++++++++++++++++++-
 2 files changed, 92 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 91fc8295d79..4eb21ea0709 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -5095,7 +5095,7 @@ async def user_get_requests():
 
 
 @router.post(
-    "/user/block",
+    "/end_user/block",
     tags=["End User Management"],
     dependencies=[Depends(user_api_key_auth)],
 )
@@ -5133,11 +5133,12 @@ async def block_user(data: BlockUsers):
 
         return {"blocked_users": records}
     except Exception as e:
+        verbose_proxy_logger.error(f"An error occurred - {str(e)}")
         raise HTTPException(status_code=500, detail={"error": str(e)})
 
 
 @router.post(
-    "/user/unblock",
+    "/end_user/unblock",
     tags=["End User Management"],
     dependencies=[Depends(user_api_key_auth)],
 )
diff --git a/litellm/tests/test_blocked_user_list.py b/litellm/tests/test_blocked_user_list.py
index b40d8296c30..d3f9f6a1a32 100644
--- a/litellm/tests/test_blocked_user_list.py
+++ b/litellm/tests/test_blocked_user_list.py
@@ -6,6 +6,7 @@ import sys, os, asyncio, time, random
 from datetime import datetime
 import traceback
 from dotenv import load_dotenv
+from fastapi import Request
 
 load_dotenv()
 import os
@@ -22,18 +23,87 @@ from litellm import Router, mock_completion
 from litellm.proxy.utils import ProxyLogging
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.caching import DualCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
+
+import pytest, logging, asyncio
+import litellm, asyncio
+from litellm.proxy.proxy_server import (
+    new_user,
+    generate_key_fn,
+    user_api_key_auth,
+    user_update,
+    delete_key_fn,
+    info_key_fn,
+    update_key_fn,
+    generate_key_fn,
+    generate_key_helper_fn,
+    spend_user_fn,
+    spend_key_fn,
+    view_spend_logs,
+    user_info,
+    block_user,
+)
+from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
+from litellm._logging import verbose_proxy_logger
+
+verbose_proxy_logger.setLevel(level=logging.DEBUG)
+
+from litellm.proxy._types import (
+    NewUserRequest,
+    GenerateKeyRequest,
+    DynamoDBArgs,
+    KeyRequest,
+    UpdateKeyRequest,
+    GenerateKeyRequest,
+    BlockUsers,
+)
+from litellm.proxy.utils import DBClient
+from starlette.datastructures import URL
+from litellm.caching import DualCache
+
+proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+
+
+@pytest.fixture
+def prisma_client():
+    from litellm.proxy.proxy_cli import append_query_params
+
+    ### add connection pool + pool timeout args
+    params = {"connection_limit": 100, "pool_timeout": 60}
+    database_url = os.getenv("DATABASE_URL")
+    modified_url = append_query_params(database_url, params)
+    os.environ["DATABASE_URL"] = modified_url
+
+    # Assuming DBClient is a class that needs to be instantiated
+    prisma_client = PrismaClient(
+        database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
+    )
+
+    # Reset litellm.proxy.proxy_server.prisma_client to None
+    litellm.proxy.proxy_server.custom_db_client = None
+    litellm.proxy.proxy_server.litellm_proxy_budget_name = (
+        f"litellm-proxy-budget-{time.time()}"
+    )
+    litellm.proxy.proxy_server.user_custom_key_generate = None
+
+    return prisma_client
 
 
 @pytest.mark.asyncio
-async def test_block_user_check():
+async def test_block_user_check(prisma_client):
     """
     - Set a blocked user as a litellm module value
     - Test to see if a call with that user id is made, an error is raised
     - Test to see if a call without that user is passes
     """
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+
     litellm.blocked_user_list = ["user_id_1"]
 
-    blocked_user_obj = _ENTERPRISE_BlockedUserList()
+    blocked_user_obj = _ENTERPRISE_BlockedUserList(
+        prisma_client=litellm.proxy.proxy_server.prisma_client
+    )
 
     _api_key = "sk-12345"
     user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
@@ -61,3 +131,20 @@ async def test_block_user_check():
         )
     except Exception as e:
         pytest.fail(f"An error occurred - {str(e)}")
+
+
+@pytest.mark.asyncio
+async def test_block_user_db_check(prisma_client):
+    """
+    - Block end user via "/user/block"
+    - Check returned value
+    """
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    _block_users = BlockUsers(user_ids=["user_id_1"])
+    result = await block_user(data=_block_users)
+    result = result["blocked_users"]
+    assert len(result) == 1
+    assert result[0].user_id == "user_id_1"
+    assert result[0].blocked == True