mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Merge pull request #13635 from BerriAI/litellm_mock_test_fix
[Proxy] LiteLLM mock test fix
This commit is contained in:
Krish Dholakia
GitHub
commit
ab8bd70c5e
3 changed files with 216 additions and 158 deletions
2
.github/workflows/test-litellm.yml
vendored
2
.github/workflows/test-litellm.yml
vendored
|
|
@ -7,7 +7,7 @@ on:
|
|||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 25
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from unittest.mock import patch
|
|||
|
||||
import litellm
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
|
||||
from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest
|
||||
|
||||
client = TestClient(app)
|
||||
|
|
@ -24,58 +25,71 @@ async def test_create_and_get_tag():
|
|||
"""
|
||||
Test creation of a new tag and retrieving its information
|
||||
"""
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
) as mock_router, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment"
|
||||
) as mock_add_tag, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks
|
||||
mock_get_tags.return_value = {}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
) as mock_router, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment"
|
||||
) as mock_add_tag, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks
|
||||
mock_get_tags.return_value = {}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
|
||||
|
||||
# Create a new tag
|
||||
tag_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag creation
|
||||
response = client.post("/tag/new", json=tag_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag created successfully"
|
||||
assert result["tag"]["name"] == "test-tag"
|
||||
assert result["tag"]["description"] == "Test tag for unit testing"
|
||||
|
||||
# Mock updated tag config for the get request
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
# Create a new tag
|
||||
tag_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
"model_info": {"model-1": "gpt-3.5-turbo"},
|
||||
}
|
||||
}
|
||||
|
||||
# Test retrieving tag info
|
||||
info_data = {"names": ["test-tag"]}
|
||||
response = client.post("/tag/info", json=info_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "test-tag" in result
|
||||
assert result["test-tag"]["description"] == "Test tag for unit testing"
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag creation
|
||||
response = client.post("/tag/new", json=tag_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag created successfully"
|
||||
assert result["tag"]["name"] == "test-tag"
|
||||
assert result["tag"]["description"] == "Test tag for unit testing"
|
||||
|
||||
# Mock updated tag config for the get request
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
"model_info": {"model-1": "gpt-3.5-turbo"},
|
||||
}
|
||||
}
|
||||
|
||||
# Test retrieving tag info
|
||||
info_data = {"names": ["test-tag"]}
|
||||
response = client.post("/tag/info", json=info_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "test-tag" in result
|
||||
assert result["test-tag"]["description"] == "Test tag for unit testing"
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -83,16 +97,26 @@ async def test_update_tag():
|
|||
"""
|
||||
Test updating an existing tag
|
||||
"""
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Original description",
|
||||
|
|
@ -101,27 +125,30 @@ async def test_update_tag():
|
|||
"updated_at": "2023-01-01T00:00:00",
|
||||
"created_by": "user-123",
|
||||
}
|
||||
}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
|
||||
}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
|
||||
|
||||
# Update tag data
|
||||
update_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Updated description",
|
||||
"models": ["model-1", "model-2"],
|
||||
}
|
||||
# Update tag data
|
||||
update_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Updated description",
|
||||
"models": ["model-1", "model-2"],
|
||||
}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag update
|
||||
response = client.post("/tag/update", json=update_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag updated successfully"
|
||||
assert result["tag"]["description"] == "Updated description"
|
||||
assert len(result["tag"]["models"]) == 2
|
||||
assert "model-2" in result["tag"]["models"]
|
||||
# Test tag update
|
||||
response = client.post("/tag/update", json=update_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag updated successfully"
|
||||
assert result["tag"]["description"] == "Updated description"
|
||||
assert len(result["tag"]["models"]) == 2
|
||||
assert "model-2" in result["tag"]["models"]
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -129,14 +156,24 @@ async def test_delete_tag():
|
|||
"""
|
||||
Test deleting a tag
|
||||
"""
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# Mock the prisma client and _get_tags_config and _save_tags_config
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
|
||||
) as mock_get_tags, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
|
||||
) as mock_save_tags:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for deletion",
|
||||
|
|
@ -145,22 +182,25 @@ async def test_delete_tag():
|
|||
"updated_at": "2023-01-01T00:00:00",
|
||||
"created_by": "user-123",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Delete tag data
|
||||
delete_data = {"name": "test-tag"}
|
||||
# Delete tag data
|
||||
delete_data = {"name": "test-tag"}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag deletion
|
||||
response = client.post("/tag/delete", json=delete_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag deleted successfully"
|
||||
# Test tag deletion
|
||||
response = client.post("/tag/delete", json=delete_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag deleted successfully"
|
||||
|
||||
# Verify _save_tags_config was called without the deleted tag
|
||||
mock_save_tags.assert_called_once()
|
||||
# Verify _save_tags_config was called without the deleted tag
|
||||
mock_save_tags.assert_called_once()
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -17,46 +17,55 @@ def test_cost_calculation_uses_debug_level(caplog):
|
|||
This ensures cost calculation details don't appear in production logs.
|
||||
Part of fix for issue #9815.
|
||||
"""
|
||||
# Create a mock completion response
|
||||
mock_response = {
|
||||
"id": "test",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "gpt-3.5-turbo",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Test response"},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30
|
||||
# Ensure verbose_logger is set to DEBUG level to capture the debug logs
|
||||
from litellm._logging import verbose_logger
|
||||
original_level = verbose_logger.level
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
|
||||
try:
|
||||
# Create a mock completion response
|
||||
mock_response = {
|
||||
"id": "test",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "gpt-3.5-turbo",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Test response"},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Test that cost calculation logs are at DEBUG level
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
try:
|
||||
cost = completion_cost(
|
||||
completion_response=mock_response,
|
||||
model="gpt-3.5-turbo"
|
||||
)
|
||||
except Exception:
|
||||
pass # Cost calculation may fail, but we're checking log levels
|
||||
|
||||
# Find the cost calculation log records
|
||||
cost_calc_records = [
|
||||
record for record in caplog.records
|
||||
if "selected model name for cost calculation" in record.message
|
||||
]
|
||||
|
||||
# Verify that cost calculation logs are at DEBUG level
|
||||
assert len(cost_calc_records) > 0, "No cost calculation logs found"
|
||||
|
||||
for record in cost_calc_records:
|
||||
assert record.levelno == logging.DEBUG, \
|
||||
f"Cost calculation log should be DEBUG level, but was {record.levelname}"
|
||||
|
||||
# Test that cost calculation logs are at DEBUG level
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
|
||||
try:
|
||||
cost = completion_cost(
|
||||
completion_response=mock_response,
|
||||
model="gpt-3.5-turbo"
|
||||
)
|
||||
except Exception:
|
||||
pass # Cost calculation may fail, but we're checking log levels
|
||||
|
||||
# Find the cost calculation log records
|
||||
cost_calc_records = [
|
||||
record for record in caplog.records
|
||||
if "selected model name for cost calculation" in record.message
|
||||
]
|
||||
|
||||
# Verify that cost calculation logs are at DEBUG level
|
||||
assert len(cost_calc_records) > 0, "No cost calculation logs found"
|
||||
|
||||
for record in cost_calc_records:
|
||||
assert record.levelno == logging.DEBUG, \
|
||||
f"Cost calculation log should be DEBUG level, but was {record.levelname}"
|
||||
finally:
|
||||
# Restore original logger level
|
||||
verbose_logger.setLevel(original_level)
|
||||
|
||||
|
||||
def test_batch_cost_calculation_uses_debug_level(caplog):
|
||||
|
|
@ -65,29 +74,38 @@ def test_batch_cost_calculation_uses_debug_level(caplog):
|
|||
"""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
from litellm.types.utils import Usage
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
# Create a mock usage object
|
||||
usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
|
||||
# Ensure verbose_logger is set to DEBUG level to capture the debug logs
|
||||
original_level = verbose_logger.level
|
||||
verbose_logger.setLevel(logging.DEBUG)
|
||||
|
||||
# Test that batch cost calculation logs are at DEBUG level
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
try:
|
||||
batch_cost_calculator(
|
||||
usage=usage,
|
||||
model="gpt-3.5-turbo",
|
||||
custom_llm_provider="openai"
|
||||
)
|
||||
except Exception:
|
||||
pass # May fail, but we're checking log levels
|
||||
|
||||
# Find batch cost calculation log records
|
||||
batch_cost_records = [
|
||||
record for record in caplog.records
|
||||
if "Calculating batch cost per token" in record.message
|
||||
]
|
||||
|
||||
# Verify logs exist and are at DEBUG level
|
||||
if batch_cost_records: # May not always log depending on the code path
|
||||
for record in batch_cost_records:
|
||||
assert record.levelno == logging.DEBUG, \
|
||||
f"Batch cost calculation log should be DEBUG level, but was {record.levelname}"
|
||||
try:
|
||||
# Create a mock usage object
|
||||
usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
|
||||
|
||||
# Test that batch cost calculation logs are at DEBUG level
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
|
||||
try:
|
||||
batch_cost_calculator(
|
||||
usage=usage,
|
||||
model="gpt-3.5-turbo",
|
||||
custom_llm_provider="openai"
|
||||
)
|
||||
except Exception:
|
||||
pass # May fail, but we're checking log levels
|
||||
|
||||
# Find batch cost calculation log records
|
||||
batch_cost_records = [
|
||||
record for record in caplog.records
|
||||
if "Calculating batch cost per token" in record.message
|
||||
]
|
||||
|
||||
# Verify logs exist and are at DEBUG level
|
||||
if batch_cost_records: # May not always log depending on the code path
|
||||
for record in batch_cost_records:
|
||||
assert record.levelno == logging.DEBUG, \
|
||||
f"Batch cost calculation log should be DEBUG level, but was {record.levelname}"
|
||||
finally:
|
||||
# Restore original logger level
|
||||
verbose_logger.setLevel(original_level)
|
||||
Loading…
Add table
Reference in a new issue