add tests

Author: Mantisus

src/crawlee/storage_clients/_redis/_dataset_client.py

@@ -126,8 +126,8 @@ async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
         if isinstance(data, dict):
             data = [data]
 
-        async with self._get_pipeline() as pipe:
-            try:
+        try:
+            async with self._get_pipeline() as pipe:
                 pipe.json().arrappend(self._items_key, '$', *data)
                 await self._update_metadata(
                     pipe,
@@ -136,8 +136,8 @@ async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
                     ),
                 )
 
-            except RedisError as e:
-                raise StorageWriteError(e) from e
+        except RedisError as e:
+            raise StorageWriteError(e) from e
 
     @override
     async def get_data(

src/crawlee/storage_clients/_redis/_key_value_store_client.py

@@ -143,9 +143,8 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
             content_type=content_type,
             size=size,
         )
-
-        async with self._get_pipeline() as pipe:
-            try:
+        try:
+            async with self._get_pipeline() as pipe:
                 # redis-py typing issue
                 await await_redis_response(pipe.hset(self._items_key, key, value_bytes))  # ty: ignore[invalid-argument-type]
 
@@ -160,8 +159,8 @@ async def set_value(self, *, key: str, value: Any, content_type: str | None = No
                     pipe, **MetadataUpdateParams(update_accessed_at=True, update_modified_at=True)
                 )
 
-            except RedisError as e:
-                raise StorageWriteError(e) from e
+        except RedisError as e:
+            raise StorageWriteError(e) from e
 
     @override
     async def get_value(self, *, key: str) -> KeyValueStoreRecord | None:

tests/unit/storage_clients/_file_system/test_fs_dataset_client.py

@@ -4,11 +4,13 @@
 import json
 from pathlib import Path
 from typing import TYPE_CHECKING
+from unittest.mock import patch
 
 import pytest
 
 from crawlee._consts import METADATA_FILENAME
 from crawlee.configuration import Configuration
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients import FileSystemStorageClient
 
 if TYPE_CHECKING:
@@ -151,3 +153,15 @@ async def test_data_persistence_across_reopens() -> None:
     assert data.items[0] == test_data
 
     await reopened_client.drop()
+
+
+async def test_error_handling_on_push_failure(dataset_client: FileSystemDatasetClient) -> None:
+    """Test that StorageWriteError is raised when file writing fails."""
+    with (
+        patch('crawlee.storage_clients._file_system._dataset_client.atomic_write', side_effect=OSError('disk full')),
+        pytest.raises(StorageWriteError) as exc_info,
+    ):
+        await dataset_client.push_data({'test': 'data'})
+
+    assert isinstance(exc_info.value.cause, OSError)
+    assert str(exc_info.value.cause) == 'disk full'

tests/unit/storage_clients/_file_system/test_fs_kvs_client.py

@@ -3,11 +3,13 @@
 import asyncio
 import json
 from typing import TYPE_CHECKING
+from unittest.mock import patch
 
 import pytest
 
 from crawlee._consts import METADATA_FILENAME
 from crawlee.configuration import Configuration
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients import FileSystemStorageClient
 
 if TYPE_CHECKING:
@@ -199,3 +201,18 @@ async def test_data_persistence_across_reopens(configuration: Configuration) ->
     assert record.value == test_value
 
     await reopened_client.drop()
+
+
+async def test_error_handling_on_set_failure(kvs_client: FileSystemKeyValueStoreClient) -> None:
+    """Test that StorageWriteError is raised when file writing fails."""
+    with (
+        patch(
+            'crawlee.storage_clients._file_system._key_value_store_client.atomic_write',
+            side_effect=OSError('disk full'),
+        ),
+        pytest.raises(StorageWriteError) as exc_info,
+    ):
+        await kvs_client.set_value(key='test', value='data')
+
+    assert isinstance(exc_info.value.cause, OSError)
+    assert str(exc_info.value.cause) == 'disk full'

tests/unit/storage_clients/_redis/test_redis_dataset_client.py

@@ -2,9 +2,12 @@
 
 import asyncio
 from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from redis.exceptions import RedisError
 
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients import RedisStorageClient
 from crawlee.storage_clients._redis._utils import await_redis_response
 
@@ -144,3 +147,22 @@ async def test_metadata_record_updates(dataset_client: RedisDatasetClient) -> No
     assert metadata.created_at == initial_created
     assert metadata.modified_at > initial_modified
     assert metadata.accessed_at > accessed_after_get
+
+
+async def test_error_handling_on_push_failure(dataset_client: RedisDatasetClient) -> None:
+    """Test that StorageWriteError is raised when Redis writing fails."""
+    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.object(dataset_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
+        pytest.raises(StorageWriteError) as exc_info,
+    ):
+        await dataset_client.push_data({'test': 'data'})
+
+    assert isinstance(exc_info.value.cause, RedisError)
+    assert str(exc_info.value.cause) == 'connection lost'

tests/unit/storage_clients/_redis/test_redis_kvs_client.py

@@ -3,9 +3,12 @@
 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.errors import StorageWriteError
 from crawlee.storage_clients import RedisStorageClient
 from crawlee.storage_clients._redis._utils import await_redis_response
 
@@ -215,3 +218,22 @@ async def test_metadata_record_updates(kvs_client: RedisKeyValueStoreClient) ->
     assert metadata.created_at == initial_created
     assert metadata.modified_at > initial_modified
     assert metadata.accessed_at > accessed_after_read
+
+
+async def test_error_handling_on_set_failure(kvs_client: RedisKeyValueStoreClient) -> None:
+    """Test that StorageWriteError is raised when Redis writing fails."""
+    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.object(kvs_client.redis, 'pipeline', return_value=mock_pipeline_ctx),
+        pytest.raises(StorageWriteError) as exc_info,
+    ):
+        await kvs_client.set_value(key='test', value='test-value')
+
+    assert isinstance(exc_info.value.cause, RedisError)
+    assert str(exc_info.value.cause) == 'connection lost'

tests/unit/storage_clients/_sql/test_sql_dataset_client.py

@@ -2,12 +2,15 @@
 
 import asyncio
 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.configuration import Configuration
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients import SqlStorageClient
 from crawlee.storage_clients._sql._db_models import DatasetItemDb, DatasetMetadataDb
 
@@ -231,3 +234,21 @@ async def test_data_persistence_across_reopens(configuration: Configuration) ->
         assert data.items[0] == test_data
 
         await reopened_client.drop()
+
+
+async def test_error_handling_on_push_failure(dataset_client: SqlDatasetClient) -> None:
+    """Test that StorageWriteError is raised when SQL writing fails."""
+    with patch(
+        'crawlee.storage_clients._sql._dataset_client.SqlDatasetClient.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 pytest.raises(StorageWriteError) as exc_info:
+            await dataset_client.push_data({'key': 'value'})
+
+        assert isinstance(exc_info.value.cause, SQLAlchemyError)
+        assert str(exc_info.value.cause) == 'db error'

tests/unit/storage_clients/_sql/test_sql_kvs_client.py

@@ -3,12 +3,15 @@
 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.configuration import Configuration
+from crawlee.errors import StorageWriteError
 from crawlee.storage_clients import SqlStorageClient
 from crawlee.storage_clients._sql._db_models import KeyValueStoreMetadataDb, KeyValueStoreRecordDb
 from crawlee.storage_clients.models import KeyValueStoreMetadata
@@ -281,3 +284,21 @@ async def test_data_persistence_across_reopens(configuration: Configuration) ->
         assert record.value == test_value
 
         await reopened_client.drop()
+
+
+async def test_error_handling_on_set_failure(kvs_client: SqlKeyValueStoreClient) -> None:
+    """Test that StorageWriteError is raised when SQL writing fails."""
+    with patch(
+        'crawlee.storage_clients._sql._key_value_store_client.SqlKeyValueStoreClient.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 pytest.raises(StorageWriteError) as exc_info:
+            await kvs_client.set_value(key='test', value='test-value')
+
+        assert isinstance(exc_info.value.cause, SQLAlchemyError)
+        assert str(exc_info.value.cause) == 'db error'