Revert "Merge pull request #2593 from BerriAI/litellm_reset_budget_fix"

This reverts commit afd363129f, reversing
changes made to c94bc94ad5.

tests/test_users.py

@@ -105,7 +105,6 @@ async def test_user_update():
     pass
 
 
-@pytest.mark.skip(reason="Flaky test on circle ci ci/cd.")
 @pytest.mark.asyncio
 async def test_users_budgets_reset():
     """
@@ -115,16 +114,6 @@ async def test_users_budgets_reset():
     - Check if value updated
     """
     get_user = f"krrish_{time.time()}@berri.ai"
-
-    async def retry_request(func, *args, _max_attempts=5, **kwargs):
-        for attempt in range(_max_attempts):
-            try:
-                return await func(*args, **kwargs)
-            except aiohttp.client_exceptions.ClientOSError as e:
-                if attempt + 1 == _max_attempts:
-                    raise  # re-raise the last ClientOSError if all attempts failed
-                print(f"Attempt {attempt+1} failed, retrying...")
-
     async with aiohttp.ClientSession() as session:
         key_gen = await new_user(
             session, 0, user_id=get_user, budget=10, budget_duration="5s"
@@ -136,12 +125,17 @@ async def test_users_budgets_reset():
         reset_at_init_value = user_info["user_info"]["budget_reset_at"]
         i = 0
         reset_at_new_value = None
-        await asyncio.sleep(70)
-        user_info = await retry_request(
-            get_user_info, session=session, get_user=get_user, call_user=key
-        )
-        reset_at_new_value = user_info["user_info"]["budget_reset_at"]
-
+        while i < 3:
+            await asyncio.sleep(70)
+            user_info = await get_user_info(
+                session=session, get_user=get_user, call_user=key
+            )
+            reset_at_new_value = user_info["user_info"]["budget_reset_at"]
+            try:
+                assert reset_at_init_value != reset_at_new_value
+                break
+            except:
+                i + 1
         assert reset_at_init_value != reset_at_new_value