Improved StepUntil

crates/viewer/re_viewer/src/viewer_test_utils/mod.rs

@@ -1,5 +1,7 @@
 use egui_kittest::Harness;
+use egui_kittest::kittest::{By, Queryable};
 use re_build_info::build_info;
+use std::sync::Arc;
 
 use crate::{
     App, AppEnvironment, AsyncRuntimeHandle, MainThreadToken, StartupOptions,
@@ -31,38 +33,89 @@ pub fn viewer_harness() -> Harness<'static, App> {
         })
 }
 
-/// Steps through the harness until the `predicate` closure returns `true`.
-pub async fn step_until<'app, 'harness, Predicate>(
-    test_description: &'static str,
-    harness: &'harness mut egui_kittest::Harness<'app, App>,
-    mut predicate: Predicate,
+/// Utility to wait until some widget appears.
+pub struct StepUntil {
     step_duration: tokio::time::Duration,
     max_duration: tokio::time::Duration,
-) where
-    Predicate: for<'a> FnMut(&'a egui_kittest::Harness<'app, App>) -> bool,
-{
-    let start_time = std::time::Instant::now();
-    let mut success = predicate(harness);
-    while !success && start_time.elapsed() <= max_duration {
-        harness.step();
-        tokio::time::sleep(step_duration).await;
-        harness.step();
-        success = predicate(harness);
+    debug_label: String,
+}
+
+impl StepUntil {
+    pub fn new(debug_label: impl ToString) -> Self {
+        Self {
+            debug_label: debug_label.to_string(),
+            step_duration: tokio::time::Duration::from_millis(100),
+            max_duration: tokio::time::Duration::from_secs(5),
+        }
+    }
+
+    /// Step duration.
+    ///
+    /// Default is 100ms.
+    pub fn step_duration(mut self, duration: tokio::time::Duration) -> Self {
+        self.step_duration = duration;
+        self
     }
 
-    if !success {
-        // Take a screenshot of the state of the harness if we failed the test.
-        // This is invaluable for debugging test failures.
-        let snapshot_path = "tests/failures";
-        harness
-            .try_snapshot_options(
-                test_description,
-                &egui_kittest::SnapshotOptions::default().output_path(snapshot_path),
-            )
-            .ok();
+    /// Max duration to wait for the predicate to be true.
+    ///
+    /// Default is 5 seconds.
+    pub fn max_duration(mut self, duration: tokio::time::Duration) -> Self {
+        self.max_duration = duration;
+        self
+    }
 
-        panic!(
-            "Timed out waiting for predicate to be true for {test_description:?}. A screenshot of the harness has been saved to `{snapshot_path}/{test_description}.new.png`."
-        );
+    /// Set the max_duration to step_duration * steps
+    pub fn steps(mut self, steps: u32) -> Self {
+        self.max_duration = self.step_duration * steps;
+        self
+    }
+
+    /// Step until the predicate returns Some.
+    pub async fn run<'app, 'harness: 'pre, 'pre, State, Predicate, R>(
+        self,
+        harness: &'harness mut egui_kittest::Harness<'app, State>,
+        mut predicate: Predicate,
+    ) -> R
+    where
+        Predicate: FnMut(&'pre Harness<'app, State>) -> Option<R>,
+    {
+        let start_time = std::time::Instant::now();
+        loop {
+            match predicate(harness) {
+                Some(result) => {
+                    return result;
+                }
+                None => {
+                    if start_time.elapsed() > self.max_duration {
+                        panic!(
+                            r#"Timed out waiting for "{}"
+
+Found nodes: {:?}"#,
+                            self.debug_label,
+                            harness.root()
+                        );
+                    }
+                    tokio::time::sleep(self.step_duration).await;
+                    harness.step();
+                }
+            }
+        }
+
+        if !success {
+            // Take a screenshot of the state of the harness if we failed the test.
+            // This is invaluable for debugging test failures.
+            let snapshot_path = "tests/failures";
+            harness
+                .try_snapshot_options(
+                    test_description,
+                    &egui_kittest::SnapshotOptions::default().output_path(snapshot_path),
+                )
+                .ok();
+
+            panic!(
+                "Timed out waiting for predicate to be true for {test_description:?}. A screenshot of the harness has been saved to `{snapshot_path}/{test_description}.new.png`."
+            );
+        }
     }
 }

crates/viewer/re_viewer/tests/app_kittest.rs

@@ -3,6 +3,7 @@
 use egui_kittest::kittest::Queryable as _;
 
 use re_viewer::viewer_test_utils;
+use re_viewer::viewer_test_utils::StepUntil;
 
 /// Navigates from welcome to settings screen and snapshots it.
 #[tokio::test]
@@ -12,19 +13,12 @@ async fn settings_screen() {
     harness.run_ok();
     harness.get_by_label_contains("Settings…").click();
     // Wait for the FFmpeg-check loading spinner to disappear.
-    viewer_test_utils::step_until(
-        "Settings screen shows up with FFMpeg binary not found error",
-        &mut harness,
-        |harness| {
-            harness
-                .query_by_label_contains(
-                    "The specified FFmpeg binary path does not exist or is not a file.",
-                )
-                .is_some()
-        },
-        tokio::time::Duration::from_millis(100),
-        tokio::time::Duration::from_secs(5),
-    )
-    .await;
+    StepUntil::new("Settings screen shows up with FFMpeg binary not found error")
+        .run(&mut harness, |harness| {
+            harness.query_by_label_contains(
+                "The specified FFmpeg binary path does not exist or is not a file.",
+            )
+        })
+        .await;
     harness.snapshot("settings_screen");
 }

tests/rust/re_integration_test/tests/re_integration.rs

@@ -2,6 +2,7 @@ use egui_kittest::SnapshotResults;
 use egui_kittest::kittest::Queryable as _;
 use re_integration_test::TestServer;
 use re_viewer::viewer_test_utils;
+use re_viewer::viewer_test_utils::StepUntil;
 
 #[tokio::test(flavor = "multi_thread")]
 pub async fn dataset_ui_test() {
@@ -33,27 +34,18 @@ pub async fn dataset_ui_test() {
     harness.get_by_label("Add").click();
     harness.run_ok();
 
-    viewer_test_utils::step_until(
-        "Redap server dataset appears",
-        &mut harness,
-        |harness| harness.query_by_label_contains("my_dataset").is_some(),
-        tokio::time::Duration::from_millis(100),
-        tokio::time::Duration::from_secs(5),
-    )
-    .await;
+    StepUntil::new("Waiting for 'my_dataset' list item")
+        .run(&mut harness, |harness| {
+            harness.query_by_label_contains("my_dataset")
+        })
+        .await;
 
     harness.get_by_label("my_dataset").click();
-    viewer_test_utils::step_until(
-        "Redap recording id appears",
-        &mut harness,
-        |harness| {
-            harness
-                .query_by_label_contains("new_recording_id")
-                .is_some()
-        },
-        tokio::time::Duration::from_millis(100),
-        tokio::time::Duration::from_secs(5),
-    )
-    .await;
+    StepUntil::new("Waiting for dataset table (new_recording_id) to load")
+        .run(&mut harness, |harness| {
+            harness.query_by_label_contains("new_recording_id")
+        })
+        .await;
+
     snapshot_results.add(harness.try_snapshot("dataset_ui_table"));
 }