test_connection_banner: assert that token/sid association removed on shutdown

Author: masenf

tests/integration/test_connection_banner.py

@@ -8,7 +8,9 @@
 
 from reflex import constants
 from reflex.environment import environment
+from reflex.istate.manager import StateManagerRedis
 from reflex.testing import AppHarness, WebDriver
+from reflex.utils.token_manager import RedisTokenManager
 
 from .utils import SessionStorage
 
@@ -127,17 +129,21 @@ def has_cloud_banner(driver: WebDriver) -> bool:
         return True
 
 
-def _assert_token(connection_banner, driver):
+def _assert_token(connection_banner, driver) -> str:
     """Poll for backend to be up.
 
     Args:
         connection_banner: AppHarness instance.
         driver: Selenium webdriver instance.
+
+    Returns:
+        The token if found, raises an assertion error otherwise.
     """
     ss = SessionStorage(driver)
     assert connection_banner._poll_for(lambda: ss.get("token") is not None), (
         "token not found"
     )
+    return ss.get("token")
 
 
 @pytest.mark.asyncio
@@ -151,9 +157,25 @@ async def test_connection_banner(connection_banner: AppHarness):
     assert connection_banner.backend is not None
     driver = connection_banner.frontend()
 
-    _assert_token(connection_banner, driver)
+    token = _assert_token(connection_banner, driver)
     AppHarness.expect(lambda: not has_error_modal(driver))
 
+    # Check that the token association was established.
+    app_event_namespace = connection_banner.app_instance.event_namespace
+    assert app_event_namespace is not None
+    app_token_manager = app_event_namespace._token_manager
+    assert app_token_manager is not None
+    assert token in app_token_manager.token_to_sid
+    sid_before = app_token_manager.token_to_sid[token]
+    if isinstance(connection_banner.state_manager, StateManagerRedis):
+        assert isinstance(app_token_manager, RedisTokenManager)
+        assert (
+            await connection_banner.state_manager.redis.get(
+                app_token_manager._get_redis_key(token)
+            )
+            == b"1"
+        )
+
     delay_button = driver.find_element(By.ID, "delay")
     increment_button = driver.find_element(By.ID, "increment")
     counter_element = driver.find_element(By.ID, "counter")
@@ -176,6 +198,17 @@ async def test_connection_banner(connection_banner: AppHarness):
     # Error modal should now be displayed
     AppHarness.expect(lambda: has_error_modal(driver))
 
+    # The token association should have been removed when the server exited.
+    assert token not in app_token_manager.token_to_sid
+    if isinstance(connection_banner.state_manager, StateManagerRedis):
+        assert isinstance(app_token_manager, RedisTokenManager)
+        assert (
+            await connection_banner.state_manager.redis.get(
+                app_token_manager._get_redis_key(token)
+            )
+            is None
+        )
+
     # Increment the counter with backend down
     increment_button.click()
     assert connection_banner.poll_for_value(counter_element, exp_not_equal="0") == "1"
@@ -189,6 +222,19 @@ async def test_connection_banner(connection_banner: AppHarness):
     # Banner should be gone now
     AppHarness.expect(lambda: not has_error_modal(driver))
 
+    # After reconnecting, the token association should be re-established.
+    if isinstance(connection_banner.state_manager, StateManagerRedis):
+        assert isinstance(app_token_manager, RedisTokenManager)
+        assert (
+            await connection_banner.state_manager.redis.get(
+                app_token_manager._get_redis_key(token)
+            )
+            == b"1"
+        )
+    # Make sure the new connection has a different websocket sid.
+    sid_after = app_token_manager.token_to_sid[token]
+    assert sid_before != sid_after
+
     # Count should have incremented after coming back up
     assert connection_banner.poll_for_value(counter_element, exp_not_equal="1") == "2"