Add click_element utility for robust element clicking in tests

tests/integration/test_dynamic_routes.py

@@ -12,7 +12,7 @@
 
 from reflex.testing import AppHarness, WebDriver
 
-from .utils import poll_assert_event_order, poll_for_navigation
+from .utils import click_element, poll_assert_event_order, poll_for_navigation
 
 
 def DynamicRoute():
@@ -235,24 +235,18 @@ def test_on_load_navigate(
         token: The token visible in the driver browser.
     """
     assert dynamic_route.app_instance is not None
-    link = driver.find_element(By.ID, "link_page_next")
-    assert link
 
     exp_order = [f"/page/[page_id]-{ix}" for ix in range(10)]
     # click the link a few times
     for ix in range(10):
         # wait for navigation, then assert on url
         with poll_for_navigation(driver):
-            link.click()
+            click_element(driver, By.ID, "link_page_next")
         assert urlsplit(driver.current_url).path == f"/page/{ix}"
 
-        link = AppHarness.poll_for_or_raise_timeout(
-            lambda: 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(
@@ -273,9 +267,8 @@ def test_on_load_navigate(
 
     # make sure internal nav still hydrates after redirect
     exp_order += ["/page/[page_id]-11"]
-    link = driver.find_element(By.ID, "link_page_next")
     with poll_for_navigation(driver):
-        link.click()
+        click_element(driver, By.ID, "link_page_next")
     poll_assert_event_order(driver, exp_order)
 
     # load same page with a query param and make sure it passes through
@@ -300,9 +293,8 @@ def test_on_load_navigate(
 
     # next/link to a 404 and ensure we still hydrate
     exp_order += ["/404-no page id"]
-    link = driver.find_element(By.ID, "link_missing")
     with poll_for_navigation(driver):
-        link.click()
+        click_element(driver, By.ID, "link_missing")
 
     # hit a page that redirects back to dynamic page
     exp_order += ["on_load_redir-{'foo': 'bar', 'page_id': '0'}", "/page/[page_id]-0"]
@@ -330,29 +322,24 @@ def test_on_load_navigate_non_dynamic(
         driver: WebDriver instance.
     """
     assert dynamic_route.app_instance is not None
-    link = driver.find_element(By.ID, "link_page_x")
-    assert link
 
     with poll_for_navigation(driver):
-        link.click()
+        click_element(driver, By.ID, "link_page_x")
     assert urlsplit(driver.current_url).path.removesuffix("/") == "/static/x"
     poll_assert_event_order(driver, ["/static/x-no page id"])
 
     # go back to the index and navigate back to the static route
-    link = driver.find_element(By.ID, "link_index")
     with poll_for_navigation(driver):
-        link.click()
+        click_element(driver, By.ID, "link_index")
     assert urlsplit(driver.current_url).path.removesuffix("/") == ""
 
-    link = driver.find_element(By.ID, "link_page_x")
     with poll_for_navigation(driver):
-        link.click()
+        click_element(driver, By.ID, "link_page_x")
     assert urlsplit(driver.current_url).path.removesuffix("/") == "/static/x"
     poll_assert_event_order(driver, ["/static/x-no page id", "/static/x-no page id"])
 
     for _ in range(3):
-        link = driver.find_element(By.ID, "link_page_x")
-        link.click()
+        click_element(driver, By.ID, "link_page_x")
         assert urlsplit(driver.current_url).path.removesuffix("/") == "/static/x"
     poll_assert_event_order(driver, ["/static/x-no page id"] * 5)
 
@@ -399,19 +386,15 @@ def assert_content(expected: str, expect_not: str):
             )
 
     assert_content("0", "")
-    next_page_link = driver.find_element(By.ID, "next-page")
-    assert next_page_link
     with poll_for_navigation(driver):
-        next_page_link.click()
+        click_element(driver, By.ID, "next-page")
     assert (
         driver.current_url.removesuffix("/")
         == f"{frontend_url.removesuffix('/')}/arg/1"
     )
     assert_content("1", "0")
-    next_page_link = driver.find_element(By.ID, "next-page")
-    assert next_page_link
     with poll_for_navigation(driver):
-        next_page_link.click()
+        click_element(driver, By.ID, "next-page")
     assert (
         driver.current_url.removesuffix("/")
         == f"{frontend_url.removesuffix('/')}/arg/2"

tests/integration/test_event_chain.py

@@ -10,6 +10,7 @@
 
 from reflex.testing import AppHarness, WebDriver
 from tests.integration.utils import (
+    click_element,
     poll_assert_event_order,
     poll_assert_relative_event_order,
 )
@@ -622,11 +623,8 @@ def test_event_chain_on_mount(
     assert event_chain.frontend_url is not None
     driver.get(event_chain.frontend_url.removesuffix("/") + uri)
 
-    unmount_button = AppHarness.poll_for_or_raise_timeout(
-        lambda: driver.find_element(By.ID, "unmount")
-    )
     assert_token(event_chain, driver)
-    unmount_button.click()
+    click_element(driver, By.ID, "unmount")
 
     poll_assert_relative_event_order(driver, expected_counts, ordering_rules)
 

tests/integration/test_login_flow.py

@@ -126,12 +126,11 @@ def test_login_flow(
 
     login_sample.poll_for_content(login_button)
     with utils.poll_for_navigation(driver):
-        login_button.click()
+        utils.click_element(driver, By.ID, "login")
     assert driver.current_url.endswith("/login")
 
-    do_it_button = driver.find_element(By.ID, "doit")
     with utils.poll_for_navigation(driver):
-        do_it_button.click()
+        utils.click_element(driver, By.ID, "doit")
     assert driver.current_url == login_sample.frontend_url
 
     def check_auth_token_header():
@@ -143,8 +142,7 @@ def check_auth_token_header():
 
     assert AppHarness.poll_for_or_raise_timeout(check_auth_token_header) == "12345"
 
-    logout_button = driver.find_element(By.ID, "logout")
-    logout_button.click()
+    utils.click_element(driver, By.ID, "logout")
 
     state_name = login_sample.get_full_state_name(["_state"])
     AppHarness.expect(

tests/integration/test_navigation.py

@@ -8,7 +8,7 @@
 
 from reflex.testing import AppHarness
 
-from .utils import SessionStorage, poll_for_navigation
+from .utils import SessionStorage, click_element, poll_for_navigation
 
 
 def NavigationApp():
@@ -69,26 +69,19 @@ def test_navigation_app(navigation_app: AppHarness):
     ss = SessionStorage(driver)
     assert AppHarness._poll_for(lambda: ss.get("token") is not None), "token not found"
 
-    internal_link = driver.find_element(By.ID, "internal")
-
     with poll_for_navigation(driver):
-        internal_link.click()
+        click_element(driver, By.ID, "internal")
     assert urlsplit(driver.current_url).path == "/internal"
     with poll_for_navigation(driver):
         driver.back()
 
-    external_link = AppHarness.poll_for_or_raise_timeout(
-        lambda: driver.find_element(By.ID, "external")
-    )
-    external2_link = driver.find_element(By.ID, "external2")
-
-    external_link.click()
+    click_element(driver, By.ID, "external")
     # Expect a new tab to open
     AppHarness.expect(lambda: len(driver.window_handles) == 2)
 
     # Switch back to the main tab
     driver.switch_to.window(driver.window_handles[0])
 
-    external2_link.click()
+    click_element(driver, By.ID, "external2")
     # Expect another new tab to open
     AppHarness.expect(lambda: len(driver.window_handles) == 3)

tests/integration/utils.py

@@ -2,15 +2,20 @@
 
 from __future__ import annotations
 
+import time
 from collections.abc import Generator, Iterator, Sequence
 from contextlib import contextmanager
 from http.client import HTTPConnection
 from urllib.parse import urlsplit
 
+from selenium.common.exceptions import (
+    NoSuchElementException,
+    StaleElementReferenceException,
+)
 from selenium.webdriver.common.by import By
 from selenium.webdriver.remote.webdriver import WebDriver
 
-from reflex.testing import AppHarness
+from reflex.testing import DEFAULT_TIMEOUT, POLL_INTERVAL, AppHarness, TimeoutType
 
 
 def request_raw(
@@ -64,6 +69,45 @@ def poll_for_navigation(
     AppHarness.expect(lambda: prev_url != driver.current_url, timeout=timeout)
 
 
+def click_element(
+    driver: WebDriver,
+    by: str,
+    value: str,
+    timeout: TimeoutType = None,
+) -> None:
+    """Locate an element and click it, tolerating the churn of a navigation.
+
+    Client-side navigation swaps the DOM after the URL changes, so an element
+    located right after navigating may not be rendered yet, or may be unmounted
+    before the click is dispatched. Both are retried until `timeout`, since both
+    are raised while locating or validating the element reference, never after
+    the click reaches the browser. Every other error, including an invalid
+    selector or an intercepted click, propagates on the first attempt.
+
+    Args:
+        driver: WebDriver instance.
+        by: Locator strategy, one of the `By` constants.
+        value: Locator value.
+        timeout: How long to keep re-locating the element.
+
+    Raises:
+        TimeoutError: if the element could not be clicked within the timeout.
+    """
+    deadline = time.monotonic() + (
+        DEFAULT_TIMEOUT if timeout is None else float(timeout)
+    )
+    while True:
+        try:
+            driver.find_element(by, value).click()
+        except (NoSuchElementException, StaleElementReferenceException) as exc:
+            if time.monotonic() >= deadline:
+                msg = f"Could not click element {by}={value!r} while polling."
+                raise TimeoutError(msg) from exc
+            time.sleep(POLL_INTERVAL)
+        else:
+            return
+
+
 def n_expected_events(exp_event_order: Sequence[str | set[str]]) -> int:
     """Calculate the number of expected events, accounting for sets in the expected order.