[REF-1988] API to Get instance of Arbitrary State class

integration/test_client_storage.py

@@ -518,17 +518,17 @@ def set_sub_sub(var: str, value: str):
     set_sub("l6", "l6 value")
     l5 = driver.find_element(By.ID, "l5")
     l6 = driver.find_element(By.ID, "l6")
+    assert AppHarness._poll_for(lambda: l6.text == "l6 value")
     assert l5.text == "l5 value"
-    assert l6.text == "l6 value"
 
     # Switch back to main window.
     driver.switch_to.window(main_tab)
 
     # The values should have updated automatically.
     l5 = driver.find_element(By.ID, "l5")
     l6 = driver.find_element(By.ID, "l6")
+    assert AppHarness._poll_for(lambda: l6.text == "l6 value")
     assert l5.text == "l5 value"
-    assert l6.text == "l6 value"
 
     # clear the cookie jar and local storage, ensure state reset to default
     driver.delete_all_cookies()

integration/test_state_inheritance.py

@@ -1,14 +1,29 @@
 """Test state inheritance."""
 
-import time
+from contextlib import suppress
 from typing import Generator
 
 import pytest
+from selenium.common.exceptions import NoAlertPresentException
+from selenium.webdriver.common.alert import Alert
 from selenium.webdriver.common.by import By
 
 from reflex.testing import DEFAULT_TIMEOUT, AppHarness, WebDriver
 
 
+def get_alert_or_none(driver: WebDriver) -> Alert | None:
+    """Switch to an alert if present.
+
+    Args:
+        driver: WebDriver instance.
+
+    Returns:
+        The alert if present, otherwise None.
+    """
+    with suppress(NoAlertPresentException):
+        return driver.switch_to.alert
+
+
 def raises_alert(driver: WebDriver, element: str) -> None:
     """Click an element and check that an alert is raised.
 
@@ -18,8 +33,8 @@ def raises_alert(driver: WebDriver, element: str) -> None:
     """
     btn = driver.find_element(By.ID, element)
     btn.click()
-    time.sleep(0.2)  # wait for the alert to appear
-    alert = driver.switch_to.alert
+    alert = AppHarness._poll_for(lambda: get_alert_or_none(driver))
+    assert isinstance(alert, Alert)
     assert alert.text == "clicked"
     alert.accept()
 
@@ -355,7 +370,7 @@ def test_state_inheritance(
     child3_other_mixin_btn = driver.find_element(By.ID, "child3-other-mixin-btn")
     child3_other_mixin_btn.click()
     child2_other_mixin_value = state_inheritance.poll_for_content(
-        child2_other_mixin, exp_not_equal="other_mixin"
+        child2_other_mixin, exp_not_equal="Child2.clicked.1"
     )
     child2_computed_mixin_value = state_inheritance.poll_for_content(
         child2_computed_other_mixin, exp_not_equal="other_mixin"

reflex/.templates/jinja/web/utils/context.js.jinja2

@@ -25,10 +25,20 @@ export const clientStorage = {}
 
 {% if state_name %}
 export const state_name = "{{state_name}}"
-export const onLoadInternalEvent = () => [
-    Event('{{state_name}}.{{const.update_vars_internal}}', {vars: hydrateClientStorage(clientStorage)}),
-    Event('{{state_name}}.{{const.on_load_internal}}')
-]
+export const onLoadInternalEvent = () => {
+    const internal_events = [];
+    const client_storage_vars = hydrateClientStorage(clientStorage);
+    if (client_storage_vars && Object.keys(client_storage_vars).length !== 0) {
+        internal_events.push(
+            Event(
+                '{{state_name}}.{{const.update_vars_internal}}',
+                {vars: client_storage_vars},
+            ),
+        );
+    }
+    internal_events.push(Event('{{state_name}}.{{const.on_load_internal}}'));
+    return internal_events;
+}
 
 export const initialEvents = () => [
     Event('{{state_name}}.{{const.hydrate}}'),

reflex/.templates/web/utils/state.js

@@ -585,7 +585,7 @@ export const useEventLoop = (
       if (storage_to_state_map[e.key]) {
         const vars = {}
         vars[storage_to_state_map[e.key]] = e.newValue
-        const event = Event(`${state_name}.update_vars_internal`, {vars: vars})
+        const event = Event(`${state_name}.update_vars_internal_state.update_vars_internal`, {vars: vars})
         addEvents([event], e);
       }
     };

reflex/app.py

@@ -69,9 +69,11 @@
     State,
     StateManager,
     StateUpdate,
+    _substate_key,
     code_uses_state_contexts,
 )
 from reflex.utils import console, exceptions, format, prerequisites, types
+from reflex.utils.exec import is_testing_env
 from reflex.utils.imports import ImportVar
 
 # Define custom types.
@@ -159,10 +161,9 @@ def __init__(self, *args, **kwargs):
             )
         super().__init__(*args, **kwargs)
         state_subclasses = BaseState.__subclasses__()
-        is_testing_env = constants.PYTEST_CURRENT_TEST in os.environ
 
         # Special case to allow test cases have multiple subclasses of rx.BaseState.
-        if not is_testing_env:
+        if not is_testing_env():
             # Only one Base State class is allowed.
             if len(state_subclasses) > 1:
                 raise ValueError(
@@ -176,7 +177,8 @@ def __init__(self, *args, **kwargs):
                     deprecation_version="0.3.5",
                     removal_version="0.5.0",
                 )
-            if len(State.class_subclasses) > 0:
+            # 2 substates are built-in and not considered when determining if app is stateless.
+            if len(State.class_subclasses) > 2:
                 self.state = State
         # Get the config
         config = get_config()
@@ -1002,7 +1004,7 @@ async def upload_file(request: Request, files: List[UploadFile]):
             )
 
         # Get the state for the session.
-        substate_token = token + "_" + handler.rpartition(".")[0]
+        substate_token = _substate_key(token, handler.rpartition(".")[0])
         state = await app.state_manager.get_state(substate_token)
 
         # get the current session ID

reflex/compiler/utils.py

@@ -138,12 +138,12 @@ def compile_state(state: Type[BaseState]) -> dict:
         A dictionary of the compiled state.
     """
     try:
-        initial_state = state().dict(initial=True)
+        initial_state = state(_reflex_internal_init=True).dict(initial=True)
     except Exception as e:
         console.warn(
             f"Failed to compile initial state with computed vars, excluding them: {e}"
         )
-        initial_state = state().dict(include_computed=False)
+        initial_state = state(_reflex_internal_init=True).dict(include_computed=False)
     return format.format_state(initial_state)
 
 

reflex/constants/compiler.py

@@ -59,9 +59,9 @@ class CompileVars(SimpleNamespace):
     # The name of the function for converting a dict to an event.
     TO_EVENT = "Event"
     # The name of the internal on_load event.
-    ON_LOAD_INTERNAL = "on_load_internal"
+    ON_LOAD_INTERNAL = "on_load_internal_state.on_load_internal"
     # The name of the internal event to update generic state vars.
-    UPDATE_VARS_INTERNAL = "update_vars_internal"
+    UPDATE_VARS_INTERNAL = "update_vars_internal_state.update_vars_internal"
 
 
 class PageNames(SimpleNamespace):