fix testing for budget manager

This commit is contained in:
Krrish Dholakia 2023-09-12 16:11:49 -07:00
parent 7b90ca4d00
commit 27f1051792
2 changed files with 4 additions and 4 deletions

View file

@ -19,7 +19,7 @@ def test_user_budget_enough():
try:
user = "1234"
# create a budget for a user
budget_manager.create_budget(total_budget=10, user=user)
budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# check if a given call can be made
data = {
@ -41,7 +41,7 @@ def test_user_budget_not_enough():
try:
user = "12345"
# create a budget for a user
budget_manager.create_budget(total_budget=0, user=user)
budget_manager.create_budget(total_budget=0, user=user, duration="daily")
# check if a given call can be made
data = {
@ -109,4 +109,4 @@ def test_reset_on_duration():
# Make sure the budget was actually reset
assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired"
except Exception as e:
pytest.fail(f"An error occurred - {str(e)}")
pytest.fail(f"An error occurred - {str(e)}")

View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "0.1.606"
version = "0.1.607"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT License"