mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
5e1c5c0bd1
* test: move offline proxy auth, guardrail hook, spend, budget reset, GCS payload and pass-through tests into tests/unit Relocate 53 legacy tests that pass with no network or keys into the tests/unit files that mirror the code they exercise, drop 2 llm_guard tests already covered by test_llm_guard_call_type_aliases, and delete the emptied legacy files. * test: freeze the reset clock and run the real reset methods in the moved budget reset tests Pin datetime in reset_budget_job and timezone_utils to a fixed instant and await the service-hook tasks the job spawns instead of sleeping. Drive the per-user failure with a malformed budget_duration row from the prisma mock so the real _reset_budget_for_key/user/team run instead of patched replacements. * test: use fixed timestamps in the moved pass-through and spend log tests * test: mock the moderation HTTP call and type the moved moderation and spend log tests Serve the OpenAI moderation response through an httpx MockTransport so the real async_make_request runs, annotate the moved tests' locals with Final, and pass get_logging_payload typed kwargs instead of an untyped input_args dict. * test: drive the moved proxy tests through boundaries and annotate their locals with Final Serve GCS downloads and the Router moderation call through httpx mocks with fake Google credentials, let the real pass-through success handler fail on the mocked upstream response, and fold the Vertex live route endpoint check into the existing route test. Annotate moved locals with Final, build the budget reset rows without rebinding, and type the remaining fakes without **kwargs
112 lines
2.9 KiB
Python
112 lines
2.9 KiB
Python
# What is this?
|
|
## This tests the llm guard integration
|
|
|
|
import asyncio
|
|
import random
|
|
|
|
# What is this?
|
|
## Unit test for presidio pii masking
|
|
import time
|
|
import traceback
|
|
from datetime import datetime
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
import pytest
|
|
from fastapi import Request, Response
|
|
from starlette.datastructures import URL
|
|
|
|
import litellm
|
|
from litellm import Router, mock_completion
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm_enterprise.enterprise_callbacks.secret_detection import (
|
|
_ENTERPRISE_SecretDetection,
|
|
)
|
|
from litellm.proxy.proxy_server import chat_completion
|
|
|
|
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
|
|
|
### UNIT TESTS FOR OpenAI Moderation ###
|
|
|
|
|
|
class testLogger(CustomLogger):
|
|
|
|
def __init__(self):
|
|
self.logged_message = None
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
print(f"On Async Success")
|
|
|
|
self.logged_message = kwargs.get("messages")
|
|
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "fake-model",
|
|
"litellm_params": {
|
|
"model": "openai/fake",
|
|
"api_base": FAKE_OPENAI_API_BASE,
|
|
"api_key": "sk-98765",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_completion_request_with_redaction():
|
|
"""
|
|
IMPORTANT Enterprise Test - Do not delete it:
|
|
Makes a /chat/completions request on LiteLLM Proxy
|
|
|
|
Ensures that the secret is redacted EVEN on the callback
|
|
"""
|
|
from litellm.proxy import proxy_server
|
|
|
|
setattr(proxy_server, "llm_router", router)
|
|
_test_logger = testLogger()
|
|
litellm.callbacks = [_ENTERPRISE_SecretDetection(), _test_logger]
|
|
litellm.turn_on_debug()
|
|
|
|
# Prepare the query string
|
|
query_params = "param1=value1¶m2=value2"
|
|
|
|
# Create the Request object with query parameters
|
|
request = Request(
|
|
scope={
|
|
"type": "http",
|
|
"method": "POST",
|
|
"path": "/chat/completions",
|
|
"headers": [(b"content-type", b"application/json")],
|
|
"query_string": query_params.encode(),
|
|
}
|
|
)
|
|
|
|
request._url = URL(url="/chat/completions")
|
|
|
|
async def return_body():
|
|
return b'{"model": "fake-model", "messages": [{"role": "user", "content": "Hello here is my OPENAI_API_KEY = sk-98765"}]}'
|
|
|
|
request.body = return_body
|
|
|
|
response = await chat_completion(
|
|
request=request,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-98765",
|
|
token="hashed_sk-98765",
|
|
),
|
|
fastapi_response=Response(),
|
|
)
|
|
|
|
await asyncio.sleep(3)
|
|
|
|
print("Info in callback after running request=", _test_logger.logged_message)
|
|
|
|
assert _test_logger.logged_message == [
|
|
{"role": "user", "content": "Hello here is my OPENAI_API_KEY = [REDACTED]"}
|
|
]
|
|
pass
|