diff --git a/tests/unit/storages/conftest.py b/tests/unit/storages/conftest.py index a28c319e0e..443061789c 100644 --- a/tests/unit/storages/conftest.py +++ b/tests/unit/storages/conftest.py @@ -14,14 +14,16 @@ ) if TYPE_CHECKING: + from collections.abc import AsyncGenerator + from fakeredis import FakeAsyncRedis @pytest.fixture(params=['memory', 'file_system', 'sql', 'redis']) -def storage_client( +async def storage_client( request: pytest.FixtureRequest, redis_client: FakeAsyncRedis, -) -> StorageClient: +) -> AsyncGenerator[StorageClient, None]: """Parameterized fixture to test with different storage clients.""" storage_client: StorageClient @@ -36,4 +38,7 @@ def storage_client( else: storage_client = FileSystemStorageClient() service_locator.set_storage_client(storage_client) - return storage_client + yield storage_client + + if isinstance(storage_client, SqlStorageClient): + await storage_client.close()