From 27f1051792176a7eb1fe3b72b72bccd6378d24e9 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 12 Sep 2023 16:11:49 -0700
Subject: [PATCH] fix testing for budget manager

---
 litellm/tests/test_budget_manager.py | 6 +++---
 pyproject.toml                       | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py
index c94e729e6..2fc1a029b 100644
--- a/litellm/tests/test_budget_manager.py
+++ b/litellm/tests/test_budget_manager.py
@@ -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)}")
\ No newline at end of file
+        pytest.fail(f"An error occurred - {str(e)}")
diff --git a/pyproject.toml b/pyproject.toml
index 3407318b3..117814b8c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"