Basic test working 🎉

Cargo.lock

@@ -5625,6 +5625,7 @@ dependencies = [
 name = "re_blueprint_tree"
 version = "0.22.0-alpha.1+dev"
 dependencies = [
+ "ahash",
  "egui",
  "itertools 0.13.0",
  "re_chunk_store",

crates/viewer/re_blueprint_tree/Cargo.toml

@@ -37,5 +37,8 @@ smallvec.workspace = true
 
 [dev-dependencies]
 re_viewer_context = { workspace = true, features = ["testing"] }
+re_viewport_blueprint = { workspace = true, features = ["testing"] }
 re_chunk_store.workspace = true
-re_view_spatial.workspace = true
+re_view_spatial.workspace = true
+
+ahash.workspace = true

crates/viewer/re_blueprint_tree/tests/blueprint_tree_tests.rs

@@ -1,19 +1,21 @@
 use egui::Vec2;
-use std::sync::Arc;
 
 use re_blueprint_tree::BlueprintTree;
 use re_chunk_store::external::re_chunk::ChunkBuilder;
-use re_chunk_store::{Chunk, RowId};
+use re_chunk_store::RowId;
 use re_entity_db::external::re_chunk_store::LatestAtQuery;
-use re_log_types::{build_frame_nr, EntityPath};
+use re_log_types::build_frame_nr;
 use re_types::archetypes::Points3D;
 use re_viewer_context::{
-    blueprint_timeline, test_context::TestContext, RecommendedView, ViewClass, ViewId,
+    blueprint_timeline, test_context::TestContext, RecommendedView, ViewClass,
+};
+use re_viewport_blueprint::{
+    test_context_ext::TestContextExt as _, ViewBlueprint, ViewportBlueprint,
 };
-use re_viewport_blueprint::{ViewBlueprint, ViewportBlueprint};
 
+/// Basic blueprint panel test
 #[test]
-fn blueprint_panel_should_match_snapshot() {
+fn basic_blueprint_panel_should_match_snapshot() {
     let mut test_context = TestContext::default();
 
     test_context.register_view_class::<re_view_spatial::SpatialView3D>();
@@ -22,27 +24,19 @@ fn blueprint_panel_should_match_snapshot() {
     test_context.log_entity("/entity/1".into(), add_point_to_chunk_builder);
     test_context.log_entity("/entity/2".into(), add_point_to_chunk_builder);
 
-    egui::__run_test_ctx(|egui_ctx| {
-        test_context.run(egui_ctx, |ctx| {
-            let viewport_blueprint = ViewportBlueprint::try_from_db(
-                &test_context.blueprint_store,
-                &test_context.blueprint_query,
-            );
-            viewport_blueprint.add_views(
-                std::iter::once(ViewBlueprint::new(
-                    re_view_spatial::SpatialView3D::identifier(),
-                    RecommendedView::root(),
-                )),
-                None,
-                None,
-            );
-
-            viewport_blueprint.save_to_blueprint_store(ctx, &ctx.view_class_registry);
-        })
+    test_context.setup_viewport_blueprint(|_, blueprint| {
+        blueprint.add_views(
+            std::iter::once(ViewBlueprint::new(
+                re_view_spatial::SpatialView3D::identifier(),
+                RecommendedView::root(),
+            )),
+            None,
+            None,
+        );
     });
 
     let blueprint_tree = BlueprintTree::default();
-    run_blueprint_panel_and_save_snapshot(test_context, blueprint_tree, "blueprint_panel")
+    run_blueprint_panel_and_save_snapshot(test_context, blueprint_tree, "basic_blueprint_panel")
 }
 
 fn add_point_to_chunk_builder(builder: ChunkBuilder) -> ChunkBuilder {

crates/viewer/re_viewer_context/src/query_context.rs

@@ -49,6 +49,7 @@ impl QueryContext<'_> {
 }
 
 /// The result of executing a single data query for a specific view.
+#[derive(Debug)]
 pub struct DataQueryResult {
     /// The [`DataResultTree`] for the query
     pub tree: DataResultTree,

crates/viewer/re_viewer_context/src/test_context.rs

@@ -214,7 +214,7 @@ impl TestContext {
         entity_path: EntityPath,
         build_chunk: impl FnOnce(ChunkBuilder) -> ChunkBuilder,
     ) {
-        let mut builder = build_chunk(Chunk::builder(entity_path));
+        let builder = build_chunk(Chunk::builder(entity_path));
         self.recording_store
             .add_chunk(&Arc::new(builder.build().unwrap()))
             .unwrap();

crates/viewer/re_viewport_blueprint/Cargo.toml

@@ -18,6 +18,10 @@ workspace = true
 [package.metadata.docs.rs]
 all-features = true
 
+[features]
+## Enable for testing utilities.
+testing = ["re_viewer_context/testing"]
+
 [dependencies]
 re_chunk.workspace = true
 re_chunk_store.workspace = true

crates/viewer/re_viewport_blueprint/src/lib.rs

@@ -5,6 +5,8 @@
 mod auto_layout;
 mod container;
 mod entity_add_info;
+#[cfg(feature = "testing")]
+pub mod test_context_ext;
 pub mod ui;
 mod view;
 mod view_contents;

crates/viewer/re_viewport_blueprint/src/test_context_ext.rs

@@ -0,0 +1,100 @@
+use ahash::HashMap;
+
+use re_viewer_context::{test_context::TestContext, Contents, ViewClassExt, ViewerContext};
+
+use crate::ViewportBlueprint;
+
+/// Extension trait to [`TestContext`] for blueprint-related features.
+pub trait TestContextExt {
+    /// See docstring on the implementation below.
+    fn setup_viewport_blueprint(
+        &mut self,
+        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint),
+    );
+}
+
+impl TestContextExt for TestContext {
+    /// Inspect or update the blueprint of a [`TestContext`].
+    ///
+    /// This helper works by deserializing the current blueprint, providing it to the provided
+    /// closure, and saving it back to the blueprint store. The closure should call the appropriate
+    /// methods of [`ViewportBlueprint`] to inspect and/or create views and containers as required.
+    ///
+    /// Each time [`setup_viewport_blueprint`], it entirely recomputes the "query results", i.e.,
+    /// the [`re_viewer_context::DataResult`]s that each view contains, based on the current content
+    /// of the recording store.
+    ///
+    /// Important pre-requisite:
+    /// - The view classes used by view must be already registered (see
+    ///   [`TestContext::register_view_class`]).
+    /// - The data store must be already populated for the views to have any content (see, e.g.,
+    ///   [`TestContext::log_entity`]).
+    ///
+    fn setup_viewport_blueprint(
+        &mut self,
+        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint),
+    ) {
+        let mut setup_blueprint: Option<_> = Some(setup_blueprint);
+
+        egui::__run_test_ctx(|egui_ctx| {
+            // We use `take` to ensure that the blueprint is setup only once, since egui forces
+            // us to a `FnMut` closure.
+            if let Some(setup_blueprint) = setup_blueprint.take() {
+                self.run(egui_ctx, |ctx| {
+                    let mut viewport_blueprint = ViewportBlueprint::try_from_db(
+                        &self.blueprint_store,
+                        &self.blueprint_query,
+                    );
+                    setup_blueprint(ctx, &mut viewport_blueprint);
+                    viewport_blueprint.save_to_blueprint_store(ctx);
+                });
+
+                self.handle_system_commands();
+
+                // Reload the blueprint store and execute all view queries.
+                let viewport_blueprint =
+                    ViewportBlueprint::try_from_db(&self.blueprint_store, &self.blueprint_query);
+
+                let applicable_entities_per_visualizer = self
+                    .view_class_registry
+                    .applicable_entities_for_visualizer_systems(&self.recording_store.store_id());
+                let mut query_results = HashMap::default();
+
+                self.run(egui_ctx, |ctx| {
+                    viewport_blueprint.visit_contents(&mut |contents, _| {
+                        if let Contents::View(view_id) = contents {
+                            let view_blueprint = viewport_blueprint.view(view_id).unwrap();
+                            let class_identifier = view_blueprint.class_identifier();
+
+                            let data_query_result = {
+                                let visualizable_entities = ctx
+                                    .view_class_registry
+                                    .get_class_or_log_error(class_identifier)
+                                    .determine_visualizable_entities(
+                                        &applicable_entities_per_visualizer,
+                                        ctx.recording(),
+                                        &ctx.view_class_registry
+                                            .new_visualizer_collection(class_identifier),
+                                        &view_blueprint.space_origin,
+                                    );
+
+                                view_blueprint.contents.execute_query(
+                                    ctx.store_context,
+                                    ctx.view_class_registry,
+                                    ctx.blueprint_query,
+                                    *view_id,
+                                    &visualizable_entities,
+                                )
+                            };
+                            query_results.insert(*view_id, data_query_result);
+                        }
+
+                        true
+                    })
+                });
+
+                self.query_results = query_results;
+            }
+        });
+    }
+}