mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
32 lines
698 B
Python
32 lines
698 B
Python
import json
|
|
import os
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from litellm.cost_calculator import response_cost_calculator
|
|
|
|
|
|
def test_cost_calculator():
|
|
class MockResponse(BaseModel):
|
|
_hidden_params = {"additional_headers": {"x-litellm-response-cost": 1000}}
|
|
|
|
result = response_cost_calculator(
|
|
response_object=MockResponse(),
|
|
model="",
|
|
custom_llm_provider=None,
|
|
call_type="",
|
|
optional_params={},
|
|
cache_hit=None,
|
|
base_model=None,
|
|
)
|
|
|
|
assert result == 1000
|