diff --git a/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py index 2bf4100b..046380c2 100644 --- a/tests/aio/query/test_query_session_pool.py +++ b/tests/aio/query/test_query_session_pool.py @@ -71,12 +71,18 @@ 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 index b22f6092..b3f18160 100644 --- a/tests/query/test_query_session_pool.py +++ b/tests/query/test_query_session_pool.py @@ -61,12 +61,18 @@ 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):