diff --git a/tests/aio/query/test_query_session_pool.py b/tests/aio/query/test_query_session_pool.py index 2bf4100b..f86ff3ed 100644 --- 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 index b22f6092..4c88ae77 100644 --- 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):