fix tests

Author: Mantisus

src/crawlee/_utils/retry.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import asyncio
+from asyncio import sleep as _retry_sleep  # Using alias for testing purposes
 from datetime import timedelta
 from functools import wraps
 from typing import TYPE_CHECKING, ParamSpec, TypeVar
@@ -24,12 +24,16 @@ def retry_on_error(
         max_attempts: Maximum number of attempts including the first one.
         base_delay: Base delay between retries; doubles on each subsequent attempt.
     """
+    if max_attempts < 1:
+        raise ValueError('max_attempts must be at least 1')
 
-    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+    if base_delay < timedelta(0):
+        raise ValueError('base_delay must be a non-negative timedelta')
 
-        if max_attempts < 1:
-            raise ValueError('max_attempts must be at least 1')
+    if not exception_types:
+        raise ValueError('At least one exception type must be specified')
 
+    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
         @wraps(func)
         async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             base_delay_seconds = base_delay.total_seconds()
@@ -39,7 +43,7 @@ async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                 except Exception as exc:  # noqa: PERF203
                     if not isinstance(exc, exception_types) or attempt >= max_attempts - 1:
                         raise
-                    await asyncio.sleep(base_delay_seconds * (2**attempt))
+                    await _retry_sleep(base_delay_seconds * (2**attempt))
             raise RuntimeError('Unreachable')
 
         return wrapper

tests/unit/_utils/test_retry.py

@@ -32,7 +32,7 @@ async def func() -> bool:
             raise ValueError('transient')
         return True
 
-    with patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+    with patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep:
         result = await func()
 
     assert result is True
@@ -46,7 +46,7 @@ async def func() -> None:
         raise ValueError('persistent')
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock),
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock),
         pytest.raises(ValueError, match='persistent'),
     ):
         await func()
@@ -61,7 +61,7 @@ async def func() -> None:
         raise TypeError('not retryable')
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         pytest.raises(TypeError),
     ):
         await func()
@@ -76,7 +76,7 @@ async def func() -> None:
         raise ValueError('test backoff')
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         pytest.raises(ValueError, match='test backoff'),
     ):
         await func()

tests/unit/storage_clients/_redis/test_redis_dataset_client.py

@@ -158,7 +158,7 @@ async def test_error_handling_on_push_failure(dataset_client: RedisDatasetClient
     mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         patch.object(dataset_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
         pytest.raises(RedisError),
     ):
@@ -178,7 +178,7 @@ async def test_push_data_does_not_retry_on_unexpected_exception(dataset_client:
     mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         patch.object(dataset_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
         pytest.raises(ValueError, match='unexpected error'),
     ):

tests/unit/storage_clients/_redis/test_redis_kvs_client.py

@@ -229,7 +229,7 @@ async def test_error_handling_on_set_failure(kvs_client: RedisKeyValueStoreClien
     mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         patch.object(kvs_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
         pytest.raises(RedisError),
     ):
@@ -249,7 +249,7 @@ async def test_set_value_does_not_retry_on_unexpected_exception(kvs_client: Redi
     mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
 
     with (
-        patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
         patch.object(kvs_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
         pytest.raises(ValueError, match='unexpected error'),
     ):

tests/unit/storage_clients/_redis/test_redis_rq_client.py

@@ -3,8 +3,10 @@
 import asyncio
 import json
 from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from redis.exceptions import RedisError
 
 from crawlee import Request
 from crawlee.storage_clients import RedisStorageClient
@@ -250,3 +252,43 @@ async def test_deduplication(rq_client: RedisRequestQueueClient) -> None:
     # Verify no more requests are available
     request3 = await rq_client.fetch_next_request()
     assert request3 is None
+
+
+async def test_error_handling_on_get_metadata_failure(rq_client: RedisRequestQueueClient) -> None:
+    """Test that get_metadata properly handles Redis errors and retries."""
+    mock_pipe = MagicMock()
+    mock_pipe.execute = AsyncMock(side_effect=RedisError('connection lost'))
+
+    mock_pipeline_ctx = MagicMock()
+    mock_pipeline_ctx.__aenter__ = AsyncMock(return_value=mock_pipe)
+    mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
+
+    with (
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
+        patch.object(rq_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
+        pytest.raises(RedisError),
+    ):
+        await rq_client.get_metadata()
+
+    # Verify that retry logic was attempted
+    assert mock_sleep.call_count == 2  # Assuming default max_attempts=3
+
+
+async def test_get_metadata_does_not_retry_on_unexpected_exception(rq_client: RedisRequestQueueClient) -> None:
+    """Test that get_metadata does not retry on unexpected exceptions."""
+    mock_pipe = MagicMock()
+    mock_pipe.execute = AsyncMock(side_effect=ValueError('unexpected error'))
+
+    mock_pipeline_ctx = MagicMock()
+    mock_pipeline_ctx.__aenter__ = AsyncMock(return_value=mock_pipe)
+    mock_pipeline_ctx.__aexit__ = AsyncMock(return_value=None)
+
+    with (
+        patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
+        patch.object(rq_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
+        pytest.raises(ValueError, match='unexpected error'),
+    ):
+        await rq_client.get_metadata()
+
+    # Verify that retry logic was not attempted
+    assert mock_sleep.call_count == 0

tests/unit/storage_clients/_sql/test_sql_dataset_client.py

@@ -247,7 +247,7 @@ async def test_error_handling_on_push_failure(dataset_client: SqlDatasetClient)
         mock_get_session.return_value = mock_session
 
         with (
-            patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
             pytest.raises(SQLAlchemyError),
         ):
             await dataset_client.push_data({'test': 'data'})
@@ -268,7 +268,7 @@ async def test_push_data_does_not_retry_on_unexpected_exception(dataset_client:
         mock_get_session.return_value = mock_session
 
         with (
-            patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
             pytest.raises(ValueError, match='unexpected error'),
         ):
             await dataset_client.push_data({'test': 'data'})

tests/unit/storage_clients/_sql/test_sql_kvs_client.py

@@ -297,7 +297,7 @@ async def test_error_handling_on_set_failure(kvs_client: SqlKeyValueStoreClient)
         mock_get_session.return_value = mock_session
 
         with (
-            patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
             pytest.raises(SQLAlchemyError),
         ):
             await kvs_client.set_value(key='test', value='test-value')
@@ -318,7 +318,7 @@ async def test_set_value_does_not_retry_on_unexpected_exception(kvs_client: SqlK
         mock_get_session.return_value = mock_session
 
         with (
-            patch('crawlee._utils.retry.asyncio.sleep', new_callable=AsyncMock) as mock_sleep,
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
             pytest.raises(ValueError, match='unexpected error'),
         ):
             await kvs_client.set_value(key='test', value='test-value')

tests/unit/storage_clients/_sql/test_sql_rq_client.py

@@ -3,9 +3,11 @@
 import asyncio
 import json
 from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from sqlalchemy import inspect, select
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from crawlee import Request
@@ -233,3 +235,45 @@ async def test_data_persistence_across_reopens(configuration: Configuration) ->
         assert {request1.url, request2.url} == {'https://example.com/1', 'https://example.com/2'}
 
         await reopened_client.drop()
+
+
+async def test_error_handling_on_get_metadata_failure(rq_client: SqlRequestQueueClient) -> None:
+    """Test that get_metadata properly handles SQL errors and retries."""
+    with patch(
+        'crawlee.storage_clients._sql._request_queue_client.SqlRequestQueueClient.get_session',
+    ) as mock_get_session:
+        mock_session = AsyncMock()
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+        mock_session.execute = AsyncMock(side_effect=SQLAlchemyError('db error'))
+        mock_get_session.return_value = mock_session
+
+        with (
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
+            pytest.raises(SQLAlchemyError),
+        ):
+            await rq_client.get_metadata()
+
+    # Verify that retry logic was attempted
+    assert mock_sleep.call_count == 2  # Assuming default max_attempts=3
+
+
+async def test_get_metadata_does_not_retry_on_unexpected_exception(rq_client: SqlRequestQueueClient) -> None:
+    """Test that get_metadata does not retry on unexpected exceptions."""
+    with patch(
+        'crawlee.storage_clients._sql._request_queue_client.SqlRequestQueueClient.get_session',
+    ) as mock_get_session:
+        mock_session = AsyncMock()
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+        mock_session.execute = AsyncMock(side_effect=ValueError('unexpected error'))
+        mock_get_session.return_value = mock_session
+
+        with (
+            patch('crawlee._utils.retry._retry_sleep', new_callable=AsyncMock) as mock_sleep,
+            pytest.raises(ValueError, match='unexpected error'),
+        ):
+            await rq_client.get_metadata()
+
+    # Verify that retry logic was not attempted
+    assert mock_sleep.call_count == 0