From 0bbae2d3d5498ccba5655eabfe99c5ded48ccef3 Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Mon, 4 Dec 2023 16:45:22 -0800
Subject: [PATCH] [REF-1404] router_data not updated when processing events
 (#2255)

---
 integration/test_dynamic_routes.py |  7 +++++++
 reflex/state.py                    |  3 ---
 tests/test_app.py                  | 11 ++++++++++-
 tests/test_state.py                |  6 +++++-
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/integration/test_dynamic_routes.py b/integration/test_dynamic_routes.py
index 782e625ee..cafb9c61b 100644
--- a/integration/test_dynamic_routes.py
+++ b/integration/test_dynamic_routes.py
@@ -41,6 +41,11 @@ def DynamicRoute():
                 id="token",
             ),
             rx.input(value=DynamicState.page_id, is_read_only=True, id="page_id"),
+            rx.input(
+                value=DynamicState.router.page.raw_path,
+                is_read_only=True,
+                id="raw_path",
+            ),
             rx.link("index", href="/", id="link_index"),
             rx.link("page_X", href="/static/x", id="link_page_x"),
             rx.link(
@@ -183,11 +188,13 @@ async def test_on_load_navigate(
 
         link = driver.find_element(By.ID, "link_page_next")
         page_id_input = driver.find_element(By.ID, "page_id")
+        raw_path_input = driver.find_element(By.ID, "raw_path")
 
         assert link
         assert page_id_input
 
         assert dynamic_route.poll_for_value(page_id_input) == str(ix)
+        assert dynamic_route.poll_for_value(raw_path_input) == f"/page/{ix}/"
     await poll_for_order(exp_order)
 
     # manually load the next page to trigger client side routing in prod mode
diff --git a/reflex/state.py b/reflex/state.py
index d6157e3cb..2c34b73fb 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -1084,9 +1084,6 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         # Get the function to process the event.
         fn = functools.partial(handler.fn, state)
 
-        # Clean the state before processing the event.
-        self._clean()
-
         # Wrap the function in a try/except block.
         try:
             # Handle async functions.
diff --git a/tests/test_app.py b/tests/test_app.py
index 6f159172a..901136b83 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -28,7 +28,7 @@ from reflex.components import Box, Component, Cond, Fragment, Text
 from reflex.event import Event
 from reflex.middleware import HydrateMiddleware
 from reflex.model import Model
-from reflex.state import BaseState, State, StateManagerRedis, StateUpdate
+from reflex.state import BaseState, RouterData, State, StateManagerRedis, StateUpdate
 from reflex.style import Style
 from reflex.utils import format
 from reflex.vars import ComputedVar
@@ -961,6 +961,14 @@ async def test_dynamic_route_var_route_change_completed_on_load(
             name=f"{state.get_full_name()}.{constants.CompileVars.ON_LOAD_INTERNAL}",
             val=exp_val,
         )
+        exp_router_data = {
+            "headers": {},
+            "ip": client_ip,
+            "sid": sid,
+            "token": token,
+            **on_load_internal.router_data,
+        }
+        exp_router = RouterData(exp_router_data)
         process_coro = process(
             app,
             event=on_load_internal,
@@ -977,6 +985,7 @@ async def test_dynamic_route_var_route_change_completed_on_load(
                     f"comp_{arg_name}": exp_val,
                     constants.CompileVars.IS_HYDRATED: False,
                     # "side_effect_counter": exp_index,
+                    "router": exp_router,
                 }
             },
             events=[
diff --git a/tests/test_state.py b/tests/test_state.py
index 28cb41106..d55d026f1 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -1837,7 +1837,9 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
     assert mock_app.event_namespace is not None
     emit_mock = mock_app.event_namespace.emit
 
-    assert json.loads(emit_mock.mock_calls[0].args[1]) == {
+    first_ws_message = json.loads(emit_mock.mock_calls[0].args[1])
+    assert first_ws_message["delta"]["background_task_state"].pop("router") is not None
+    assert first_ws_message == {
         "delta": {
             "background_task_state": {
                 "order": ["background_task:start"],
@@ -2405,6 +2407,7 @@ async def test_preprocess(app_module_mock, token, test_state, expected, mocker):
         assert isinstance(update, StateUpdate)
         updates.append(update)
     assert len(updates) == 1
+    assert updates[0].delta["state"].pop("router") is not None
     assert updates[0].delta == exp_is_hydrated(state, False)
 
     events = updates[0].events
@@ -2446,6 +2449,7 @@ async def test_preprocess_multiple_load_events(app_module_mock, token, mocker):
         assert isinstance(update, StateUpdate)
         updates.append(update)
     assert len(updates) == 1
+    assert updates[0].delta["state"].pop("router") is not None
     assert updates[0].delta == exp_is_hydrated(state, False)
 
     events = updates[0].events