test: closing redis ressources

docs/usage/stores.rst

@@ -259,3 +259,11 @@ with minimal boilerplate:
 
 Without any extra configuration, every call to ``app.stores.get`` with a unique name will return a namespace for this
 name only, while re-using the underlying Redis instance.
+
+
+Store lifetime
+++++++++++++++
+
+Stores may not be automatically closed when the application is shut down.
+This is the case in particular for the RedisStore if you are not using the class method :meth:`RedisStore.with_client <.redis.RedisStore.with_client>` and passing in your own Redis instance.
+In this case you're responsible to close the Redis instance yourself.

litestar/app.py

@@ -388,7 +388,12 @@ def __init__(
 
         self._openapi_schema: OpenAPI | None = None
         self._debug: bool = True
+        self.stores: StoreRegistry = (
+            config.stores if isinstance(config.stores, StoreRegistry) else StoreRegistry(config.stores)
+        )
         self._lifespan_managers = config.lifespan
+        for store in self.stores._stores.values():
+            self._lifespan_managers.append(store)
         self._server_lifespan_managers = [p.server_lifespan for p in config.plugins or [] if isinstance(p, CLIPlugin)]
         self.experimental_features = frozenset(config.experimental_features or [])
         self.get_logger: GetLogger = get_logger_placeholder
@@ -471,10 +476,6 @@ def __init__(
 
         self.asgi_handler = self._create_asgi_handler()
 
-        self.stores: StoreRegistry = (
-            config.stores if isinstance(config.stores, StoreRegistry) else StoreRegistry(config.stores)
-        )
-
     @property
     @deprecated(version="2.6.0", kind="property", info="Use create_static_files router instead")
     def static_files_config(self) -> list[StaticFilesConfig]:

litestar/stores/base.py

@@ -9,6 +9,8 @@
 from msgspec.msgpack import encode as msgpack_encode
 
 if TYPE_CHECKING:
+    from types import TracebackType
+
     from typing_extensions import Self
 
 
@@ -76,6 +78,17 @@ async def expires_in(self, key: str) -> int | None:
         """
         raise NotImplementedError
 
+    async def __aenter__(self) -> None:  # noqa: B027
+        pass
+
+    async def __aexit__(  # noqa: B027
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        pass
+
 
 class NamespacedStore(Store):
     """A subclass of :class:`Store`, offering hierarchical namespacing.

litestar/stores/redis.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 from redis.asyncio import Redis
 from redis.asyncio.connection import ConnectionPool
@@ -14,23 +14,30 @@
 
 __all__ = ("RedisStore",)
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 
 class RedisStore(NamespacedStore):
     """Redis based, thread and process safe asynchronous key/value store."""
 
     __slots__ = ("_redis",)
 
-    def __init__(self, redis: Redis, namespace: str | None | EmptyType = Empty) -> None:
+    def __init__(
+        self, redis: Redis, namespace: str | None | EmptyType = Empty, handle_client_shutdown: bool = False
+    ) -> None:
         """Initialize :class:`RedisStore`
 
         Args:
             redis: An :class:`redis.asyncio.Redis` instance
             namespace: A key prefix to simulate a namespace in redis. If not given,
                 defaults to ``LITESTAR``. Namespacing can be explicitly disabled by passing
                 ``None``. This will make :meth:`.delete_all` unavailable.
+            handle_client_shutdown: If ``True``, handle the shutdown of the `redis` instance automatically during the store's lifespan. Should be set to `True` unless the shutdown is handled externally
         """
         self._redis = redis
         self.namespace: str | None = value_or_default(namespace, "LITESTAR")
+        self.handle_client_shutdown = handle_client_shutdown
 
         # script to get and renew a key in one atomic step
         self._get_and_renew_script = self._redis.register_script(
@@ -64,6 +71,18 @@ def __init__(self, redis: Redis, namespace: str | None | EmptyType = Empty) -> N
         """
         )
 
+    async def _shutdown(self) -> None:
+        if self.handle_client_shutdown:
+            await self._redis.aclose(close_connection_pool=True)  # type: ignore[attr-defined]
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self._shutdown()
+
     @classmethod
     def with_client(
         cls,
@@ -93,14 +112,22 @@ def with_client(
             username=username,
             password=password,
         )
-        return cls(redis=Redis(connection_pool=pool), namespace=namespace)
+        return cls(
+            redis=Redis(connection_pool=pool),
+            namespace=namespace,
+            handle_client_shutdown=True,
+        )
 
     def with_namespace(self, namespace: str) -> RedisStore:
         """Return a new :class:`RedisStore` with a nested virtual key namespace.
         The current instances namespace will serve as a prefix for the namespace, so it
         can be considered the parent namespace.
         """
-        return type(self)(redis=self._redis, namespace=f"{self.namespace}_{namespace}" if self.namespace else namespace)
+        return type(self)(
+            redis=self._redis,
+            namespace=f"{self.namespace}_{namespace}" if self.namespace else namespace,
+            handle_client_shutdown=self.handle_client_shutdown,
+        )
 
     def _make_key(self, key: str) -> str:
         prefix = f"{self.namespace}:" if self.namespace else ""

tests/conftest.py

@@ -81,6 +81,13 @@ def redis_store(redis_client: AsyncRedis) -> RedisStore:
     return RedisStore(redis=redis_client)
 
 
+@pytest.fixture()
+async def redis_store_with_client() -> RedisStore:
+    store = RedisStore.with_client(url="redis://localhost:6397")
+    await store._redis.flushall()
+    return store
+
+
 @pytest.fixture()
 def memory_store() -> MemoryStore:
     return MemoryStore()
@@ -92,10 +99,15 @@ def file_store(tmp_path: Path) -> FileStore:
 
 
 @pytest.fixture(
-    params=[pytest.param("redis_store", marks=pytest.mark.xdist_group("redis")), "memory_store", "file_store"]
+    params=[
+        lazy_fixture("redis_store"),
+        lazy_fixture("memory_store"),
+        lazy_fixture("file_store"),
+        lazy_fixture("redis_store_with_client"),
+    ]
 )
 def store(request: FixtureRequest) -> Store:
-    return cast("Store", request.getfixturevalue(request.param))
+    return cast("Store", request.param)
 
 
 @pytest.fixture

tests/e2e/test_response_caching.py

@@ -7,6 +7,7 @@
 
 import msgspec
 import pytest
+from pytest_lazyfixture import lazy_fixture
 
 from litestar import Litestar, Request, Response, get, post
 from litestar.config.compression import CompressionConfig
@@ -168,12 +169,20 @@ def handler() -> str:
     assert await app.stores.get("some_store").exists("GET/")
 
 
-async def test_with_stores(store: Store, mock: MagicMock) -> None:
+@pytest.mark.parametrize(
+    "store",
+    [
+        lazy_fixture("memory_store"),
+        lazy_fixture("file_store"),
+        lazy_fixture("redis_store"),
+    ],
+)
+async def test_with_stores(store: Store, mock: MagicMock, redis_service: None) -> None:
     @get(cache=True)
     def handler() -> str:
         return mock()  # type: ignore[no-any-return]
 
-    app = Litestar([handler], stores={"response_cache": store})
+    app = Litestar([handler], stores={"response_cache": store}, debug=True)
 
     with TestClient(app=app) as client:
         response_one = client.get("/")

tests/unit/test_stores.py

@@ -366,3 +366,16 @@ async def test_file_store_handle_rename_fail(file_store: FileStore, mocker: Mock
     await file_store.set("foo", "bar")
     mock_unlink.assert_called_once()
     assert Path(mock_unlink.call_args_list[0].args[0]).with_suffix("") == file_store.path.joinpath("foo")
+
+
+async def test_redis_store_with_client_shutdown(redis_service: None) -> None:
+    redis_store = RedisStore.with_client(url="redis://localhost:6397")
+    assert await redis_store._redis.ping()
+    # remove the private shutdown and the assert below fails
+    # the check on connection is a mimic of https://github.com/redis/redis-py/blob/d529c2ad8d2cf4dcfb41bfd93ea68cfefd81aa66/tests/test_asyncio/test_connection_pool.py#L35-L39
+    await redis_store._shutdown()
+    assert not any(
+        x.is_connected
+        for x in redis_store._redis.connection_pool._available_connections
+        + list(redis_store._redis.connection_pool._in_use_connections)
+    )

tests/unit/test_testing/test_test_client.py

@@ -2,9 +2,11 @@
 from typing import TYPE_CHECKING, Callable, Dict, NoReturn, Optional, Union, cast
 
 from _pytest.fixtures import FixtureRequest
+from pytest_lazyfixture import lazy_fixture
 
 from litestar import Controller, WebSocket, delete, head, patch, put, websocket
 from litestar.status_codes import HTTP_200_OK, HTTP_201_CREATED, HTTP_204_NO_CONTENT
+from litestar.stores.base import Store
 from litestar.testing import AsyncTestClient, WebSocketTestSession, create_test_client
 
 if TYPE_CHECKING:
@@ -23,7 +25,6 @@
 import pytest
 
 from litestar import Litestar, Request, get, post
-from litestar.stores.base import Store
 from litestar.testing import TestClient
 from litestar.utils.helpers import get_exception_group
 from tests.helpers import maybe_async, maybe_async_cm
@@ -67,6 +68,14 @@ def get_session_data(request: Request) -> Dict[str, Any]:
         assert session_data == (await maybe_async(client.get("/test"))).json()  # type: ignore[attr-defined]
 
 
+@pytest.mark.parametrize(
+    "store",
+    [
+        lazy_fixture("memory_store"),
+        lazy_fixture("file_store"),
+        lazy_fixture("redis_store"),
+    ],
+)
 @pytest.mark.parametrize("with_domain", [False, True])
 async def test_test_client_get_session_data(
     with_domain: bool,