tests/local_testing/test_budget_manager.py

# #### What this tests ####
# #    This tests calling batch_completions by running 100 messages together

# import sys, os, json
# import traceback
# import pytest

# sys.path.insert(
#     0, os.path.abspath("../..")
# )  # Adds the parent directory to the system path
# import litellm
# litellm.set_verbose = True
# from litellm import completion, BudgetManager

# budget_manager = BudgetManager(project_name="test_project", client_type="hosted")

# ## Scenario 1: User budget enough to make call
# def test_user_budget_enough():
#     try:
#         user = "1234"
#         # create a budget for a user
#         budget_manager.create_budget(total_budget=10, user=user, duration="daily")

#         # check if a given call can be made
#         data = {
#             "model": "gpt-3.5-turbo",
#             "messages": [{"role": "user", "content": "Hey, how's it going?"}]
#         }
#         if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
#             response = completion(**data)
#             print(budget_manager.update_cost(completion_obj=response, user=user))
#         else:
#             response = "Sorry - no budget!"

#         print(f"response: {response}")
#     except Exception as e:
#         pytest.fail(f"An error occurred - {str(e)}")

# ## Scenario 2: User budget not enough to make call
# def test_user_budget_not_enough():
#     try:
#         user = "12345"
#         # create a budget for a user
#         budget_manager.create_budget(total_budget=0, user=user, duration="daily")

#         # check if a given call can be made
#         data = {
#             "model": "gpt-3.5-turbo",
#             "messages": [{"role": "user", "content": "Hey, how's it going?"}]
#         }
#         model = data["model"]
#         messages = data["messages"]
#         if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
#             response = completion(**data)
#             print(budget_manager.update_cost(completion_obj=response, user=user))
#         else:
#             response = "Sorry - no budget!"

#         print(f"response: {response}")
#     except Exception:
#         pytest.fail(f"An error occurred")

# ## Scenario 3: Saving budget to client
# def test_save_user_budget():
#     try:
#         response = budget_manager.save_data()
#         if response["status"] == "error":
#             raise Exception(f"An error occurred - {json.dumps(response)}")
#         print(response)
#     except Exception as e:
#         pytest.fail(f"An error occurred: {str(e)}")

# test_save_user_budget()
# ## Scenario 4: Getting list of users
# def test_get_users():
#     try:
#         response = budget_manager.get_users()
#         print(response)
#     except Exception:
#         pytest.fail(f"An error occurred")


# ## Scenario 5: Reset budget at the end of duration
# def test_reset_on_duration():
#     try:
#         # First, set a short duration budget for a user
#         user = "123456"
#         budget_manager.create_budget(total_budget=10, user=user, duration="daily")

#         # Use some of the budget
#         data = {
#             "model": "gpt-3.5-turbo",
#             "messages": [{"role": "user", "content": "Hello!"}]
#         }
#         if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user):
#             response = litellm.completion(**data)
#             print(budget_manager.update_cost(completion_obj=response, user=user))

#         assert budget_manager.get_current_cost(user) > 0, f"Test setup failed: Budget did not decrease after completion"

#         # Now, we need to simulate the passing of time. Since we don't want our tests to actually take days, we're going
#         # to cheat a little -- we'll manually adjust the "created_at" time so it seems like a day has passed.
#         # In a real-world testing scenario, we might instead use something like the `freezegun` library to mock the system time.
#         one_day_in_seconds = 24 * 60 * 60
#         budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds

#         # Now the duration should have expired, so our budget should reset
#         budget_manager.update_budget_all_users()

#         # 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)}")

# ## Scenario 6: passing in text:
# def test_input_text_on_completion():
#     try:
#         user = "12345"
#         budget_manager.create_budget(total_budget=10, user=user, duration="daily")

#         input_text = "hello world"
#         output_text = "it's a sunny day in san francisco"
#         model = "gpt-3.5-turbo"

#         budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text)
#         print(budget_manager.get_current_cost(user))
#     except Exception as e:
#         pytest.fail(f"An error occurred - {str(e)}")

# test_input_text_on_completion()