From 9f369ec1a712d97fc11c3eeb779741b85b9dc7b1 Mon Sep 17 00:00:00 2001
From: Morty Space <morty.space@gmail.com>
Date: Wed, 18 May 2022 11:38:51 +0200
Subject: [PATCH] test: added retries for trade tests

---
 tests/test_private.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/test_private.py b/tests/test_private.py
index 9c5c48c..61bd88a 100644
--- a/tests/test_private.py
+++ b/tests/test_private.py
@@ -1,4 +1,5 @@
 import asyncio
+import functools
 import time
 
 import async_timeout
@@ -12,6 +13,26 @@ from cryptocom.exchange.structs import (
 )
 
 
+def retry(times: int):
+    """Retry function for unstable tests."""
+
+    def decorator(f):
+        @functools.wraps(f)
+        async def wrapper(*args, **kwargs):
+            nonlocal times
+            while True:
+                try:
+                    return await f(*args, **kwargs)
+                except Exception as exc:
+                    times -= 1
+                    if not times:
+                        raise exc
+
+        return wrapper
+
+    return decorator
+
+
 @pytest.mark.asyncio
 async def test_account_get_balance(account: cro.Account):
     balances = await account.get_balance()
@@ -63,6 +84,7 @@ async def test_deposit_withdrawal_history(
 
 
 @pytest.mark.asyncio
+@retry(5)
 async def test_no_duplicate_mass_limit_orders(
     exchange: cro.Exchange, account: cro.Account
 ):
@@ -93,6 +115,7 @@ async def test_no_duplicate_mass_limit_orders(
 
 
 @pytest.mark.asyncio
+@retry(5)
 async def test_account_limit_orders(
     account: cro.Account, exchange: cro.Exchange
 ):
@@ -156,6 +179,7 @@ async def listen_orders(account: cro.Account, orders):
 
 
 @pytest.mark.asyncio
+@retry(5)
 async def test_account_market_orders(
     account: cro.Account, exchange: cro.Exchange
 ):