Added LiteLLM to the stack

2025-08-18 09:40:50 +00:00
parent 0648c1968c
commit d220b04e32
2682 changed files with 533609 additions and 1 deletions
--- a/Development/litellm/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
+++ b/Development/litellm/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
@@ -0,0 +1,547 @@
+import json
+import os
+import sys
+import uuid
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.proxy._types import (
+    LiteLLM_TeamMembership,
+    LiteLLM_UserTable,
+    Member,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.management_helpers.utils import add_new_member
+
+
+@pytest.mark.asyncio
+async def test_add_new_member_uses_default_team_budget_id():
+    """
+    Test that add_new_member uses the default_team_budget_id when max_budget_in_team is None.
+
+    This test verifies that:
+    1. When max_budget_in_team is None
+    2. And default_team_budget_id is provided
+    3. The team membership is created with the default_team_budget_id
+    """
+    from litellm.proxy._types import LitellmUserRoles
+
+    # Setup test data
+    test_user_id = "test_user_123"
+    test_team_id = "test_team_456"
+    test_default_budget_id = "default_budget_789"
+    test_admin_name = "test_admin"
+
+    # Create a Member object with user_id
+    new_member = Member(user_id=test_user_id, role="user")
+
+    # Create UserAPIKeyAuth object
+    user_api_key_dict = UserAPIKeyAuth(
+        user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Mock the user table upsert operation
+    mock_user_response = MagicMock()
+    mock_user_response.model_dump.return_value = {
+        "user_id": test_user_id,
+        "user_email": None,
+        "teams": [test_team_id],
+        "user_role": "internal_user",
+    }
+    mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
+        return_value=mock_user_response
+    )
+
+    # Mock the team membership creation
+    mock_team_membership_response = MagicMock()
+    mock_team_membership_response.model_dump.return_value = {
+        "team_id": test_team_id,
+        "user_id": test_user_id,
+        "budget_id": test_default_budget_id,
+        "litellm_budget_table": None,
+    }
+    mock_prisma_client.db.litellm_teammembership.create = AsyncMock(
+        return_value=mock_team_membership_response
+    )
+
+    # Call the function with max_budget_in_team=None and a default_team_budget_id
+    result_user, result_team_membership = await add_new_member(
+        new_member=new_member,
+        max_budget_in_team=None,  # This is the key - no max budget specified
+        prisma_client=mock_prisma_client,
+        team_id=test_team_id,
+        user_api_key_dict=user_api_key_dict,
+        litellm_proxy_admin_name=test_admin_name,
+        default_team_budget_id=test_default_budget_id,  # This should be used
+    )
+
+    # Verify that the user was created/updated correctly
+    assert result_user is not None
+    assert result_user.user_id == test_user_id
+
+    # Verify that the team membership was created correctly
+    assert result_team_membership is not None
+    assert result_team_membership.team_id == test_team_id
+    assert result_team_membership.user_id == test_user_id
+    assert result_team_membership.budget_id == test_default_budget_id
+
+    # Verify that the prisma client methods were called correctly
+    mock_prisma_client.db.litellm_usertable.upsert.assert_called_once()
+    mock_prisma_client.db.litellm_teammembership.create.assert_called_once()
+
+    # Verify that no budget table creation was called (since max_budget_in_team is None)
+    assert (
+        not hasattr(mock_prisma_client.db, "litellm_budgettable")
+        or not mock_prisma_client.db.litellm_budgettable.create.called
+    )
+
+    # Verify the team membership was created with the correct budget_id
+    team_membership_call_args = (
+        mock_prisma_client.db.litellm_teammembership.create.call_args
+    )
+    assert team_membership_call_args is not None
+    create_data = team_membership_call_args.kwargs["data"]
+    assert create_data["budget_id"] == test_default_budget_id
+
+
+@pytest.mark.asyncio
+async def test_add_new_member_creates_new_budget_when_max_budget_provided():
+    """
+    Test that add_new_member creates a new budget when max_budget_in_team is provided.
+
+    This test verifies that:
+    1. When max_budget_in_team is provided
+    2. A new budget is created in the litellm_budgettable
+    3. The new budget_id is used for the team membership
+    """
+    from litellm.proxy._types import LitellmUserRoles
+
+    # Setup test data
+    test_user_id = "test_user_123"
+    test_team_id = "test_team_456"
+    test_max_budget = 100.0
+    test_new_budget_id = "new_budget_789"
+    test_admin_name = "test_admin"
+
+    # Create a Member object with user_id
+    new_member = Member(user_id=test_user_id, role="user")
+
+    # Create UserAPIKeyAuth object
+    user_api_key_dict = UserAPIKeyAuth(
+        user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Mock the user table upsert operation
+    mock_user_response = MagicMock()
+    mock_user_response.model_dump.return_value = {
+        "user_id": test_user_id,
+        "user_email": None,
+        "teams": [test_team_id],
+        "user_role": "internal_user",
+    }
+    mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(
+        return_value=mock_user_response
+    )
+
+    # Mock the budget table creation
+    mock_budget_response = MagicMock()
+    mock_budget_response.budget_id = test_new_budget_id
+    mock_prisma_client.db.litellm_budgettable.create = AsyncMock(
+        return_value=mock_budget_response
+    )
+
+    # Mock the team membership creation
+    mock_team_membership_response = MagicMock()
+    mock_team_membership_response.model_dump.return_value = {
+        "team_id": test_team_id,
+        "user_id": test_user_id,
+        "budget_id": test_new_budget_id,
+        "litellm_budget_table": None,
+    }
+    mock_prisma_client.db.litellm_teammembership.create = AsyncMock(
+        return_value=mock_team_membership_response
+    )
+
+    # Call the function with max_budget_in_team provided
+    result_user, result_team_membership = await add_new_member(
+        new_member=new_member,
+        max_budget_in_team=test_max_budget,  # This should trigger budget creation
+        prisma_client=mock_prisma_client,
+        team_id=test_team_id,
+        user_api_key_dict=user_api_key_dict,
+        litellm_proxy_admin_name=test_admin_name,
+        default_team_budget_id=None,  # Should be ignored since max_budget_in_team is provided
+    )
+
+    # Verify that the budget was created
+    mock_prisma_client.db.litellm_budgettable.create.assert_called_once()
+    budget_call_args = mock_prisma_client.db.litellm_budgettable.create.call_args
+    budget_data = budget_call_args.kwargs["data"]
+    assert budget_data["max_budget"] == test_max_budget
+    assert budget_data["created_by"] == user_api_key_dict.user_id
+    assert budget_data["updated_by"] == user_api_key_dict.user_id
+
+    # Verify that the team membership was created with the new budget_id
+    assert result_team_membership is not None
+    assert result_team_membership.budget_id == test_new_budget_id
+
+    # Verify the team membership was created with the correct budget_id
+    team_membership_call_args = (
+        mock_prisma_client.db.litellm_teammembership.create.call_args
+    )
+    assert team_membership_call_args is not None
+    create_data = team_membership_call_args.kwargs["data"]
+    assert create_data["budget_id"] == test_new_budget_id
+
+
+@pytest.mark.asyncio
+async def test_add_new_member_with_user_email():
+    """
+    Test add_new_member with user_email instead of user_id and default budget.
+
+    This test verifies that:
+    1. When new_member has user_email instead of user_id
+    2. And max_budget_in_team is None
+    3. The default_team_budget_id is used correctly
+    """
+    from litellm.proxy._types import LitellmUserRoles
+
+    # Setup test data
+    test_user_email = "test@example.com"
+    test_team_id = "test_team_456"
+    test_default_budget_id = "default_budget_789"
+    test_admin_name = "test_admin"
+
+    # Create a Member object with user_email
+    new_member = Member(user_email=test_user_email, role="user")
+
+    # Create UserAPIKeyAuth object
+    user_api_key_dict = UserAPIKeyAuth(
+        user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Mock get_data to return empty list (no existing user)
+    mock_prisma_client.get_data = AsyncMock(return_value=[])
+
+    # Mock insert_data for new user creation
+    mock_user_response = MagicMock()
+    mock_user_response.model_dump.return_value = {
+        "user_id": "generated_user_id",
+        "user_email": test_user_email,
+        "teams": [test_team_id],
+        "user_role": "internal_user",
+    }
+    mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response)
+
+    # Mock the team membership creation
+    mock_team_membership_response = MagicMock()
+    mock_team_membership_response.model_dump.return_value = {
+        "team_id": test_team_id,
+        "user_id": "generated_user_id",
+        "budget_id": test_default_budget_id,
+        "litellm_budget_table": None,
+    }
+    mock_prisma_client.db.litellm_teammembership.create = AsyncMock(
+        return_value=mock_team_membership_response
+    )
+
+    # Call the function
+    result_user, result_team_membership = await add_new_member(
+        new_member=new_member,
+        max_budget_in_team=None,
+        prisma_client=mock_prisma_client,
+        team_id=test_team_id,
+        user_api_key_dict=user_api_key_dict,
+        litellm_proxy_admin_name=test_admin_name,
+        default_team_budget_id=test_default_budget_id,
+    )
+
+    # Verify that the user was created correctly
+    assert result_user is not None
+    assert result_user.user_email == test_user_email
+
+    # Verify that the team membership was created with the default budget_id
+    assert result_team_membership is not None
+    assert result_team_membership.budget_id == test_default_budget_id
+
+    # Verify that get_data was called to check for existing user
+    mock_prisma_client.get_data.assert_called_once_with(
+        key_val={"user_email": test_user_email},
+        table_name="user",
+        query_type="find_all",
+    )
+
+    # Verify that insert_data was called to create new user
+    mock_prisma_client.insert_data.assert_called_once()
+    insert_call_args = mock_prisma_client.insert_data.call_args
+    insert_data = insert_call_args.kwargs["data"]
+    assert insert_data["user_email"] == test_user_email
+    assert insert_data["teams"] == [test_team_id]
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_with_object_permission_id():
+    """
+    Test that attach_object_permission_to_dict correctly attaches object_permission
+    when object_permission_id is present and found in database.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data
+    test_object_permission_id = "test_perm_123"
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "object_permission_id": test_object_permission_id,
+        "other_field": "other_value"
+    }
+    
+    expected_object_permission = {
+        "object_permission_id": test_object_permission_id,
+        "vector_stores": ["store1", "store2"],
+        "assistants": ["assistant1"],
+        "models": ["gpt-4", "claude-3"]
+    }
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+    
+    # Mock the object permission response
+    mock_object_permission = MagicMock()
+    mock_object_permission.model_dump.return_value = expected_object_permission
+    
+    # Mock the database query
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
+        return_value=mock_object_permission
+    )
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result
+    assert result is not None
+    assert result["user_id"] == "test_user_456"
+    assert result["object_permission_id"] == test_object_permission_id
+    assert result["other_field"] == "other_value"
+    assert result["object_permission"] == expected_object_permission
+
+    # Verify the database query was called correctly
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
+        where={"object_permission_id": test_object_permission_id}
+    )
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_without_object_permission_id():
+    """
+    Test that attach_object_permission_to_dict returns the original dict unchanged
+    when object_permission_id is not present.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data without object_permission_id
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "other_field": "other_value"
+    }
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result is unchanged
+    assert result == test_data_dict
+    assert "object_permission" not in result
+
+    # Verify no database query was made
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_object_permission_not_found():
+    """
+    Test that attach_object_permission_to_dict returns the original dict unchanged
+    when object_permission_id is present but not found in database.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data
+    test_object_permission_id = "test_perm_123"
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "object_permission_id": test_object_permission_id,
+        "other_field": "other_value"
+    }
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+    
+    # Mock the database query to return None (not found)
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
+        return_value=None
+    )
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result is unchanged
+    assert result == test_data_dict
+    assert "object_permission" not in result
+
+    # Verify the database query was called
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
+        where={"object_permission_id": test_object_permission_id}
+    )
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_with_dict_method():
+    """
+    Test that attach_object_permission_to_dict handles object permissions that use .dict() method
+    instead of .model_dump() method.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data
+    test_object_permission_id = "test_perm_123"
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "object_permission_id": test_object_permission_id,
+        "other_field": "other_value"
+    }
+    
+    expected_object_permission = {
+        "object_permission_id": test_object_permission_id,
+        "vector_stores": ["store1"],
+        "assistants": []
+    }
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+    
+    # Mock the object permission response that uses .dict() method
+    mock_object_permission = MagicMock()
+    mock_object_permission.model_dump.side_effect = AttributeError("No model_dump method")
+    mock_object_permission.dict.return_value = expected_object_permission
+    
+    # Mock the database query
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
+        return_value=mock_object_permission
+    )
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result
+    assert result is not None
+    assert result["object_permission"] == expected_object_permission
+
+    # Verify both methods were attempted
+    mock_object_permission.model_dump.assert_called_once()
+    mock_object_permission.dict.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_with_none_prisma_client():
+    """
+    Test that attach_object_permission_to_dict raises ValueError when prisma_client is None.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "object_permission_id": "test_perm_123"
+    }
+
+    # Call the function with None prisma_client
+    with pytest.raises(ValueError, match="Prisma client not found"):
+        await attach_object_permission_to_dict(
+            data_dict=test_data_dict,
+            prisma_client=None  # type: ignore
+        )
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_with_empty_dict():
+    """
+    Test that attach_object_permission_to_dict handles empty dictionaries correctly.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup empty test data
+    test_data_dict = {}
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result is unchanged
+    assert result == {}
+    assert "object_permission" not in result
+
+    # Verify no database query was made
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_attach_object_permission_to_dict_with_none_object_permission_id():
+    """
+    Test that attach_object_permission_to_dict handles None object_permission_id correctly.
+    """
+    from litellm.proxy.management_helpers.object_permission_utils import attach_object_permission_to_dict
+
+    # Setup test data with None object_permission_id
+    test_data_dict = {
+        "user_id": "test_user_456",
+        "object_permission_id": None,
+        "other_field": "other_value"
+    }
+
+    # Mock the prisma client
+    mock_prisma_client = AsyncMock()
+
+    # Call the function
+    result = await attach_object_permission_to_dict(
+        data_dict=test_data_dict,
+        prisma_client=mock_prisma_client
+    )
+
+    # Verify the result is unchanged
+    assert result == test_data_dict
+    assert "object_permission" not in result
+
+    # Verify no database query was made
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_not_called()