add retry logic to normal case tests

ydb-platform · Jan 23, 2025 · ba5c216 · ba5c216
1 parent 030275b
commit ba5c216
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py
@@ -71,12 +71,19 @@ async def callee(session: QuerySession):
     )
     @pytest.mark.asyncio
     async def test_retry_tx_normal(self, pool: QuerySessionPool, tx_mode: Optional[ydb.BaseQueryTxMode]):
+        retry_no = 0
+
         async def callee(tx: QueryTxContext):
+            nonlocal retry_no
+            if retry_no < 2:
+                retry_no += 1
+                raise ydb.Unavailable("Fake fast backoff error")
             result_stream = await tx.execute("SELECT 1")
             return [result_set async for result_set in result_stream]
 
         result = await pool.retry_tx_async(callee=callee, tx_mode=tx_mode)
         assert len(result) == 1
+        assert retry_no == 2
 
     @pytest.mark.asyncio
     async def test_retry_tx_raises(self, pool: QuerySessionPool):

diff --git a/tests/query/test_query_session_pool.py b/tests/query/test_query_session_pool.py
@@ -61,12 +61,19 @@ def callee(session: QuerySession):
         ],
     )
     def test_retry_tx_normal(self, pool: QuerySessionPool, tx_mode: Optional[ydb.BaseQueryTxMode]):
+        retry_no = 0
+
         def callee(tx: QueryTxContext):
+            nonlocal retry_no
+            if retry_no < 2:
+                retry_no += 1
+                raise ydb.Unavailable("Fake fast backoff error")
             result_stream = tx.execute("SELECT 1")
             return [result_set for result_set in result_stream]
 
         result = pool.retry_tx_sync(callee=callee, tx_mode=tx_mode)
         assert len(result) == 1
+        assert retry_no == 2
 
     def test_retry_tx_raises(self, pool: QuerySessionPool):
         class CustomException(Exception):