Add disconnect/reconnect test to test_background_task

Author: masenf

tests/integration/test_background_task.py

@@ -109,6 +109,22 @@ async def yield_in_async_with_self(self):
                 yield
                 self.counter += 1
 
+        @rx.event
+        async def disconnect_reconnect(self):
+            self.counter += 1
+            yield rx.call_script("socket.disconnect()")
+            await asyncio.sleep(0.5)
+            self.counter += 1
+
+        @rx.event(background=True)
+        async def disconnect_reconnect_background(self):
+            async with self:
+                self.counter += 1
+            yield rx.call_script("socket.disconnect()")
+            await asyncio.sleep(0.5)
+            async with self:
+                self.counter += 1
+
     class OtherState(rx.State):
         @rx.event(background=True)
         async def get_other_state(self):
@@ -134,6 +150,9 @@ def index() -> rx.Component:
             rx.input(
                 id="token", value=State.router.session.client_token, is_read_only=True
             ),
+            rx.input(
+                id="sid", value=State.router.session.session_id, is_read_only=True
+            ),
             rx.hstack(
                 rx.heading(State.counter, id="counter"),
                 rx.text(State.counter_async_cv, size="1", id="counter-async-cv"),
@@ -185,6 +204,16 @@ def index() -> rx.Component:
                 on_click=State.yield_in_async_with_self,
                 id="yield-in-async-with-self",
             ),
+            rx.button(
+                "Disconnect / Reconnect",
+                on_click=State.disconnect_reconnect_background,
+                id="disconnect-reconnect",
+            ),
+            rx.button(
+                "Disconnect / Reconnect Background",
+                on_click=State.disconnect_reconnect_background,
+                id="disconnect-reconnect-background",
+            ),
             rx.button("Reset", on_click=State.reset_counter, id="reset"),
         )
 
@@ -395,3 +424,43 @@ def test_yield_in_async_with_self(
 
     yield_in_async_with_self_button.click()
     AppHarness.expect(lambda: counter.text == "2", timeout=5)
+
+
+@pytest.mark.parametrize(
+    "button_id",
+    [
+        "disconnect-reconnect",
+        "disconnect-reconnect-background",
+    ],
+)
+def test_disconnect_reconnect(
+    background_task: AppHarness,
+    driver: WebDriver,
+    token: str,
+    button_id: str,
+):
+    """Test that disconnecting and reconnecting works as expected.
+
+    Args:
+        background_task: harness for BackgroundTask app.
+        driver: WebDriver instance.
+        token: The token for the connected client.
+        button_id: The ID of the button to click.
+    """
+    counter = driver.find_element(By.ID, "counter")
+    button = driver.find_element(By.ID, button_id)
+    increment_button = driver.find_element(By.ID, "increment")
+    sid_input = driver.find_element(By.ID, "sid")
+    sid = background_task.poll_for_value(sid_input, timeout=5)
+    assert sid is not None
+
+    AppHarness.expect(lambda: counter.text == "0", timeout=5)
+    button.click()
+    AppHarness.expect(lambda: counter.text == "1", timeout=5)
+    increment_button.click()
+    # should get a new sid after the reconnect
+    assert (
+        background_task.poll_for_value(sid_input, timeout=5, exp_not_equal=sid) != sid
+    )
+    # Final update should come through on the new websocket connection
+    AppHarness.expect(lambda: counter.text == "3", timeout=5)