45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import pytest
|
|
|
|
from crawlee import service_locator
|
|
from crawlee.storage_clients import (
|
|
FileSystemStorageClient,
|
|
MemoryStorageClient,
|
|
RedisStorageClient,
|
|
SqlStorageClient,
|
|
StorageClient,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncGenerator
|
|
|
|
from fakeredis import FakeAsyncRedis
|
|
|
|
|
|
@pytest.fixture(params=['memory', 'file_system', 'sql', 'redis'])
|
|
async def storage_client(
|
|
request: pytest.FixtureRequest,
|
|
redis_client: FakeAsyncRedis,
|
|
) -> AsyncGenerator[StorageClient, None]:
|
|
"""Parameterized fixture to test with different storage clients."""
|
|
storage_client: StorageClient
|
|
|
|
storage_type = request.param
|
|
|
|
if storage_type == 'memory':
|
|
storage_client = MemoryStorageClient()
|
|
elif storage_type == 'sql':
|
|
storage_client = SqlStorageClient()
|
|
elif storage_type == 'redis':
|
|
storage_client = RedisStorageClient(redis=redis_client)
|
|
else:
|
|
storage_client = FileSystemStorageClient()
|
|
service_locator.set_storage_client(storage_client)
|
|
yield storage_client
|
|
|
|
if isinstance(storage_client, SqlStorageClient):
|
|
await storage_client.close()
|