fix AppHarness reloading (#2916)

* move AppHarness tests to module scope

* fix AppHarness reloading

* add test

* docstrings and formatting

* fix benchmarks not reloading state module

benchmarks/test_benchmark_compile_pages.py

@@ -320,6 +320,7 @@ def test_app_1_compile_time_cold(benchmark, app_with_one_page):
             app_with_one_page.app_instance.compile_()
 
     benchmark.pedantic(benchmark_fn, setup=setup, rounds=5)
+    app_with_one_page._reload_state_module()
 
 
 @pytest.mark.benchmark(
@@ -345,6 +346,7 @@ def test_app_1_compile_time_warm(benchmark, app_with_one_page):
             app_with_one_page.app_instance.compile_()
 
     benchmark(benchmark_fn)
+    app_with_one_page._reload_state_module()
 
 
 @pytest.mark.skipif(constants.IS_WINDOWS, reason=WINDOWS_SKIP_REASON)
@@ -373,6 +375,7 @@ def test_app_10_compile_time_cold(benchmark, app_with_ten_pages):
             app_with_ten_pages.app_instance.compile_()
 
     benchmark.pedantic(benchmark_fn, setup=setup, rounds=5)
+    app_with_ten_pages._reload_state_module()
 
 
 @pytest.mark.benchmark(
@@ -398,6 +401,7 @@ def test_app_10_compile_time_warm(benchmark, app_with_ten_pages):
             app_with_ten_pages.app_instance.compile_()
 
     benchmark(benchmark_fn)
+    app_with_ten_pages._reload_state_module()
 
 
 @pytest.mark.skipif(constants.IS_WINDOWS, reason=WINDOWS_SKIP_REASON)
@@ -426,6 +430,7 @@ def test_app_100_compile_time_cold(benchmark, app_with_hundred_pages):
             app_with_hundred_pages.app_instance.compile_()
 
     benchmark.pedantic(benchmark_fn, setup=setup, rounds=5)
+    app_with_hundred_pages._reload_state_module()
 
 
 @pytest.mark.benchmark(
@@ -451,6 +456,7 @@ def test_app_100_compile_time_warm(benchmark, app_with_hundred_pages):
             app_with_hundred_pages.app_instance.compile_()
 
     benchmark(benchmark_fn)
+    app_with_hundred_pages._reload_state_module()
 
 
 @pytest.mark.skipif(constants.IS_WINDOWS, reason=WINDOWS_SKIP_REASON)
@@ -479,6 +485,7 @@ def test_app_1000_compile_time_cold(benchmark, app_with_thousand_pages):
             app_with_thousand_pages.app_instance.compile_()
 
     benchmark.pedantic(benchmark_fn, setup=setup, rounds=5)
+    app_with_thousand_pages._reload_state_module()
 
 
 @pytest.mark.benchmark(
@@ -504,6 +511,7 @@ def test_app_1000_compile_time_warm(benchmark, app_with_thousand_pages):
             app_with_thousand_pages.app_instance.compile_()
 
     benchmark(benchmark_fn)
+    app_with_thousand_pages._reload_state_module()
 
 
 @pytest.mark.skip
@@ -532,6 +540,7 @@ def test_app_10000_compile_time_cold(benchmark, app_with_ten_thousand_pages):
             app_with_ten_thousand_pages.app_instance.compile_()
 
     benchmark.pedantic(benchmark_fn, setup=setup, rounds=5)
+    app_with_ten_thousand_pages._reload_state_module()
 
 
 @pytest.mark.skip
@@ -555,3 +564,4 @@ def test_app_10000_compile_time_warm(benchmark, app_with_ten_thousand_pages):
             app_with_ten_thousand_pages.app_instance.compile_()
 
     benchmark(benchmark_fn)
+    app_with_ten_thousand_pages._reload_state_module()

integration/shared/state.py

@@ -0,0 +1,8 @@
+"""Simple module which contains one reusable reflex state class."""
+import reflex as rx
+
+
+class SharedState(rx.State):
+    """Shared state class for reflexers using librarys."""
+
+    pass

integration/test_background_task.py

@@ -97,7 +97,7 @@ def BackgroundTask():
     app.add_page(index)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def background_task(
     tmp_path_factory,
 ) -> Generator[AppHarness, None, None]:

integration/test_call_script.py

@@ -230,7 +230,7 @@ def CallScript():
         )
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def call_script(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start CallScript app at tmp_path via AppHarness.
 

integration/test_client_storage.py

@@ -111,7 +111,7 @@ def ClientSide():
     app.add_page(index, route="/foo")
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def client_side(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start ClientSide app at tmp_path via AppHarness.
 

integration/test_dynamic_routes.py

@@ -74,7 +74,7 @@ def DynamicRoute():
     app.add_custom_404_page(on_load=DynamicState.on_load)  # type: ignore
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def dynamic_route(
     app_harness_env: Type[AppHarness], tmp_path_factory
 ) -> Generator[AppHarness, None, None]:

integration/test_event_actions.py

@@ -137,7 +137,7 @@ def TestEventAction():
     app.add_page(index)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def event_action(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start TestEventAction app at tmp_path via AppHarness.
 

integration/test_event_chain.py

@@ -247,7 +247,7 @@ def EventChain():
     app.add_page(on_mount_yield_chain)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def event_chain(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start EventChain app at tmp_path via AppHarness.
 

integration/test_form_submit.py

@@ -141,7 +141,7 @@ def FormSubmitName(form_component):
 
 
 @pytest.fixture(
-    scope="session",
+    scope="module",
     params=[
         functools.partial(FormSubmit, form_component="rx.form.root"),
         functools.partial(FormSubmitName, form_component="rx.form.root"),

integration/test_login_flow.py

@@ -47,7 +47,7 @@ def LoginSample():
     app.add_page(login)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def login_sample(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start LoginSample app at tmp_path via AppHarness.
 

integration/test_server_side_event.py

@@ -76,7 +76,7 @@ def ServerSideEvent():
         )
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def server_side_event(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start ServerSideEvent app at tmp_path via AppHarness.
 

integration/test_shared_state.py

@@ -0,0 +1,76 @@
+"""Test shared state."""
+from __future__ import annotations
+
+from typing import Generator
+
+import pytest
+
+from reflex.testing import AppHarness, WebDriver
+
+
+def SharedStateApp():
+    """Test that shared state works as expected."""
+    import reflex as rx
+    from integration.shared.state import SharedState
+
+    class State(SharedState):
+        pass
+
+    def index() -> rx.Component:
+        return rx.vstack()
+
+    app = rx.App()
+    app.add_page(index)
+
+
+@pytest.fixture
+def shared_state(
+    tmp_path_factory,
+) -> Generator[AppHarness, None, None]:
+    """Start SharedStateApp at tmp_path via AppHarness.
+
+    Args:
+        tmp_path_factory: pytest tmp_path_factory fixture
+
+    Yields:
+        running AppHarness instance
+
+    """
+    with AppHarness.create(
+        root=tmp_path_factory.mktemp("shared_state"),
+        app_source=SharedStateApp,  # type: ignore
+    ) as harness:
+        yield harness
+
+
+@pytest.fixture
+def driver(shared_state: AppHarness) -> Generator[WebDriver, None, None]:
+    """Get an instance of the browser open to the shared_state app.
+
+    Args:
+        shared_state: harness for SharedStateApp
+
+    Yields:
+        WebDriver instance.
+
+    """
+    assert shared_state.app_instance is not None, "app is not running"
+    driver = shared_state.frontend()
+    try:
+        yield driver
+    finally:
+        driver.quit()
+
+
+def test_shared_state(
+    shared_state: AppHarness,
+    driver: WebDriver,
+):
+    """Test that 2 AppHarness instances can share a state (f.e. from a library).
+
+    Args:
+        shared_state: harness for SharedStateApp.
+        driver: WebDriver instance.
+
+    """
+    assert shared_state.app_instance is not None

integration/test_state_inheritance.py

@@ -200,7 +200,7 @@ def StateInheritance():
     app.add_page(index)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def state_inheritance(
     tmp_path_factory,
 ) -> Generator[AppHarness, None, None]:

integration/test_upload.py

@@ -119,7 +119,7 @@ def UploadFile():
     app.add_page(index)
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def upload_file(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start UploadFile app at tmp_path via AppHarness.
 

integration/test_var_operations.py

@@ -586,7 +586,7 @@ def VarOperations():
         )
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="module")
 def var_operations(tmp_path_factory) -> Generator[AppHarness, None, None]:
     """Start VarOperations app at tmp_path via AppHarness.
 

reflex/state.py

@@ -2976,8 +2976,12 @@ def reload_state_module(
     Args:
         module: The module to reload.
         state: Recursive argument for the state class to reload.
+
     """
     for subclass in tuple(state.class_subclasses):
         reload_state_module(module=module, state=subclass)
         if subclass.__module__ == module and module is not None:
             state.class_subclasses.remove(subclass)
+            state._always_dirty_substates.discard(subclass.get_name())
+    state._init_var_dependency_dicts()
+    state.get_class_substate.cache_clear()

reflex/testing.py

@@ -40,7 +40,13 @@ import reflex.utils.build
 import reflex.utils.exec
 import reflex.utils.prerequisites
 import reflex.utils.processes
-from reflex.state import BaseState, State, StateManagerMemory, StateManagerRedis
+from reflex.state import (
+    BaseState,
+    State,
+    StateManagerMemory,
+    StateManagerRedis,
+    reload_state_module,
+)
 
 try:
     from selenium import webdriver  # pyright: ignore [reportMissingImports]
@@ -73,10 +79,6 @@ else:
     FRONTEND_POPEN_ARGS["start_new_session"] = True
 
 
-# Save a copy of internal substates to reset after each test.
-INTERNAL_STATES = State.class_subclasses.copy()
-
-
 # borrowed from py3.11
 class chdir(contextlib.AbstractContextManager):
     """Non thread-safe context manager to change the current working directory."""
@@ -229,11 +231,6 @@ class AppHarness:
             reflex.config.get_config(reload=True)
             # Clean out any `rx.page` decorators from other tests.
             reflex.app.DECORATED_PAGES.clear()
-            # reset rx.State subclasses
-            State.class_subclasses.clear()
-            State.class_subclasses.update(INTERNAL_STATES)
-            State._always_dirty_substates = set()
-            State.get_class_substate.cache_clear()
             # Ensure the AppHarness test does not skip State assignment due to running via pytest
             os.environ.pop(reflex.constants.PYTEST_CURRENT_TEST, None)
             self.app_module = reflex.utils.prerequisites.get_compiled_app(reload=True)
@@ -244,6 +241,10 @@ class AppHarness:
         else:
             self.state_manager = self.app_instance._state_manager
 
+    def _reload_state_module(self):
+        """Reload the rx.State module to avoid conflict when reloading."""
+        reload_state_module(module=f"{self.app_name}.{self.app_name}")
+
     def _get_backend_shutdown_handler(self):
         if self.backend is None:
             raise RuntimeError("Backend was not initialized.")
@@ -361,6 +362,8 @@ class AppHarness:
 
     def stop(self) -> None:
         """Stop the frontend and backend servers."""
+        self._reload_state_module()
+
         if self.backend is not None:
             self.backend.should_exit = True
         if self.frontend_process is not None: