From 9140453d0c724663cdff072e05c3f2a07609736f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 19 Mar 2024 19:47:07 -0700
Subject: [PATCH] test: make test call more reliable

---
 tests/test_users.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tests/test_users.py b/tests/test_users.py
index 57f24e3a0d4..636cb0078d5 100644
--- a/tests/test_users.py
+++ b/tests/test_users.py
@@ -114,6 +114,16 @@ 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"
@@ -126,8 +136,8 @@ async def test_users_budgets_reset():
         i = 0
         reset_at_new_value = None
         while i < 3:
-            user_info = await get_user_info(
-                session=session, get_user=get_user, call_user=key
+            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"]
             try: