mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
e9d40a8f73
A name bound twice keeps only the second binding. In `tests/` that is nearly always a repeated import, harmless but misleading, and the same rule is what catches the cases that are not harmless: a local that shadows an import the module still calls, and a second `def test_x` that quietly replaces the first. 311 of the 344 sites were repeated imports and came out with ruff's own fix. The remaining 33 needed a decision. Four modules imported a name they never used because a local definition below already shadowed it. Two comprehensions bound `call` over `unittest.mock.call`, which those modules import and use. One test rebound the two module handles its nested reload closure had captured. One class attribute shadowed an unused `status` import. The load-test fixtures move to a conftest, which is how pytest is meant to share them, so the test module no longer imports three fixture names it never calls. The nine `prisma_client` parameters keep a narrow `noqa`: pytest resolves that fixture by name before the body runs, so the parameter never shadows anything.
77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
# What is this?
|
|
## This tests the braintrust integration
|
|
|
|
import asyncio
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from datetime import datetime
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi import Request
|
|
|
|
load_dotenv()
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import logging
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
|
|
|
|
|
def test_braintrust_logging():
|
|
|
|
litellm.set_verbose = True
|
|
|
|
http_client = HTTPHandler()
|
|
|
|
with patch(
|
|
"litellm.integrations.braintrust_logging.HTTPHandler.post",
|
|
new=MagicMock(),
|
|
) as mock_client:
|
|
# set braintrust as a callback, litellm will send the data to braintrust
|
|
litellm.callbacks = ["braintrust"]
|
|
|
|
# openai call
|
|
response = litellm.completion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
|
|
)
|
|
|
|
time.sleep(2)
|
|
mock_client.assert_called()
|
|
|
|
|
|
def test_braintrust_logging_specific_project_id():
|
|
|
|
litellm.set_verbose = True
|
|
|
|
with patch(
|
|
"litellm.integrations.braintrust_logging.HTTPHandler.post",
|
|
new=MagicMock(),
|
|
) as mock_client:
|
|
# set braintrust as a callback, litellm will send the data to braintrust
|
|
litellm.callbacks = ["braintrust"]
|
|
|
|
response = litellm.completion(
|
|
model="openai/gpt-4o",
|
|
messages=[{"content": "Hello, how are you?", "role": "user"}],
|
|
metadata={"project_id": "123"},
|
|
)
|
|
|
|
time.sleep(2)
|
|
|
|
# Check that the log was inserted into the correct project
|
|
mock_client.assert_called()
|
|
_, kwargs = mock_client.call_args
|
|
assert "url" in kwargs
|
|
assert (
|
|
kwargs["url"] == "https://api.braintrustdata.com/v1/project_logs/123/insert"
|
|
)
|