AppHarness: call state_manager.close() for all state managers

Author: masenf

reflex/testing.py

@@ -305,8 +305,9 @@ def _get_backend_shutdown_handler(self):
 
         async def _shutdown(*args, **kwargs) -> None:
             # ensure redis is closed before event loop
-            if self.app_instance is not None and isinstance(
-                self.app_instance._state_manager, StateManagerRedis
+            if (
+                self.app_instance is not None
+                and self.app_instance._state_manager is not None
             ):
                 with contextlib.suppress(ValueError):
                     await self.app_instance._state_manager.close()
@@ -358,6 +359,12 @@ async def _reset_backend_state_manager(self):
         Raises:
             RuntimeError: when the state manager cannot be reset
         """
+        if (
+            self.app_instance is not None
+            and self.app_instance._state_manager is not None
+        ):
+            with contextlib.suppress(RuntimeError):
+                await self.app_instance._state_manager.close()
         if (
             self.app_instance is not None
             and isinstance(
@@ -366,8 +373,6 @@ async def _reset_backend_state_manager(self):
             )
             and self.app_instance._state is not None
         ):
-            with contextlib.suppress(RuntimeError):
-                await self.app_instance._state_manager.close()
             self.app_instance._state_manager = StateManagerRedis.create(
                 state=self.app_instance._state,
             )
@@ -716,8 +721,7 @@ async def get_state(self, token: str) -> BaseState:
         try:
             return await self.state_manager.get_state(token)
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                await self.state_manager.close()
+            await self.state_manager.close()
 
     async def set_state(self, token: str, **kwargs) -> None:
         """Set the state associated with the given token.
@@ -738,8 +742,7 @@ async def set_state(self, token: str, **kwargs) -> None:
         try:
             await self.state_manager.set_state(token, state)
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                await self.state_manager.close()
+            await self.state_manager.close()
 
     @contextlib.asynccontextmanager
     async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
@@ -769,9 +772,8 @@ async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
             async with self.app_instance.modify_state(token) as state:
                 yield state
         finally:
-            if isinstance(self.state_manager, StateManagerRedis):
-                self.app_instance._state_manager = app_state_manager
-                await self.state_manager.close()
+            self.app_instance._state_manager = app_state_manager
+            await self.state_manager.close()
 
     def token_manager(self) -> TokenManager:
         """Get the token manager for the app instance.