simple pytest benchmark for measuring event <=> state update round trip time (#2489)

integration/test_large_state.py

@@ -0,0 +1,90 @@
+"""Test large state."""
+import time
+
+import jinja2
+import pytest
+from selenium.webdriver.common.by import By
+
+from reflex.testing import AppHarness, WebDriver
+
+LARGE_STATE_APP_TEMPLATE = """
+import reflex as rx
+
+class State(rx.State):
+    var0: int = 0
+    {% for i in range(1, var_count) %}
+    var{{ i }}: str = "{{ i }}" * 10000
+    {% endfor %}
+
+    def increment_var0(self):
+        self.var0 += 1
+
+
+def index() -> rx.Component:
+    return rx.box(rx.button(State.var0, on_click=State.increment_var0, id="button"))
+
+app = rx.App()
+app.add_page(index)
+"""
+
+
+def get_driver(large_state) -> WebDriver:
+    """Get an instance of the browser open to the large_state app.
+
+    Args:
+        large_state: harness for LargeState app
+
+    Returns:
+        WebDriver instance.
+    """
+    assert large_state.app_instance is not None, "app is not running"
+    return large_state.frontend()
+
+
+@pytest.mark.parametrize("var_count", [1, 10, 100, 1000, 10000])
+def test_large_state(var_count: int, tmp_path_factory, benchmark):
+    """Measure how long it takes for button click => state update to round trip.
+
+    Args:
+        var_count: number of variables to store in the state
+        tmp_path_factory: pytest fixture
+        benchmark: pytest fixture
+
+    Raises:
+        TimeoutError: if the state doesn't update within 30 seconds
+    """
+    template = jinja2.Template(LARGE_STATE_APP_TEMPLATE)
+    large_state_rendered = template.render(var_count=var_count)
+
+    with AppHarness.create(
+        root=tmp_path_factory.mktemp(f"large_state"),
+        app_source=large_state_rendered,
+        app_name="large_state",
+    ) as large_state:
+        driver = get_driver(large_state)
+        try:
+            assert large_state.app_instance is not None
+            button = driver.find_element(By.ID, "button")
+
+            t = time.time()
+            while button.text != "0":
+                time.sleep(0.1)
+                if time.time() - t > 30.0:
+                    raise TimeoutError("Timeout waiting for initial state")
+
+            times_clicked = 0
+
+            def round_trip(clicks: int, timeout: float):
+                t = time.time()
+                for _ in range(clicks):
+                    button.click()
+                nonlocal times_clicked
+                times_clicked += clicks
+                while button.text != str(times_clicked):
+                    time.sleep(0.005)
+                    if time.time() - t > timeout:
+                        raise TimeoutError("Timeout waiting for state update")
+
+            benchmark(round_trip, clicks=10, timeout=30.0)
+        finally:
+            driver.quit()

reflex/testing.py

@@ -102,7 +102,7 @@ class AppHarness:
     """AppHarness executes a reflex app in-process for testing."""
 
     app_name: str
-    app_source: Optional[types.FunctionType | types.ModuleType]
+    app_source: Optional[types.FunctionType | types.ModuleType] | str
     app_path: pathlib.Path
     app_module_path: pathlib.Path
     app_module: Optional[types.ModuleType] = None
@@ -119,7 +119,7 @@ class AppHarness:
     def create(
         cls,
         root: pathlib.Path,
-        app_source: Optional[types.FunctionType | types.ModuleType] = None,
+        app_source: Optional[types.FunctionType | types.ModuleType | str] = None,
         app_name: Optional[str] = None,
     ) -> "AppHarness":
         """Create an AppHarness instance at root.
@@ -127,10 +127,13 @@ class AppHarness:
         Args:
             root: the directory that will contain the app under test.
             app_source: if specified, the source code from this function or module is used
-                as the main module for the app. If unspecified, then root must already
+                as the main module for the app. It may also be the raw source code text, as a str.
-                contain a working reflex app and will be used directly.
+                If unspecified, then root must already contain a working reflex app and will be used directly.
             app_name: provide the name of the app, otherwise will be derived from app_source or root.
 
+        Raises:
+            ValueError: when app_source is a string and app_name is not provided.
+
         Returns:
             AppHarness instance
         """
@@ -139,6 +142,10 @@ class AppHarness:
                 app_name = root.name.lower()
             elif isinstance(app_source, functools.partial):
                 app_name = app_source.func.__name__.lower()
+            elif isinstance(app_source, str):
+                raise ValueError(
+                    "app_name must be provided when app_source is a string."
+                )
             else:
                 app_name = app_source.__name__.lower()
         return cls(
@@ -170,16 +177,18 @@ class AppHarness:
         glbs.update(overrides)
         return glbs
 
-    def _get_source_from_func(self, func: Any) -> str:
+    def _get_source_from_app_source(self, app_source: Any) -> str:
-        """Get the source from a function or module object.
+        """Get the source from app_source.
 
         Args:
-            func: function or module object
+            app_source: function or module or str
 
         Returns:
             source code
         """
-        source = inspect.getsource(func)
+        if isinstance(app_source, str):
+            return app_source
+        source = inspect.getsource(app_source)
         source = re.sub(r"^\s*def\s+\w+\s*\(.*?\):", "", source, flags=re.DOTALL)
         return textwrap.dedent(source)
 
@@ -194,7 +203,7 @@ class AppHarness:
             source_code = "\n".join(
                 [
                     "\n".join(f"{k} = {v!r}" for k, v in app_globals.items()),
-                    self._get_source_from_func(self.app_source),
+                    self._get_source_from_app_source(self.app_source),
                 ]
             )
             with chdir(self.app_path):