Use test context in component override tests (#6499)

### What This refactor was originally a required part of something else I had to discard again. Salvaged these improvements and usages to the TestContext since since I found it independently useful and more concise than before. ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested the web demo (if applicable): * Using examples from latest `main` build: [rerun.io/viewer](https://rerun.io/viewer/pr/6499?manifest_url=https://app.rerun.io/version/main/examples_manifest.json) * Using full set of examples from `nightly` build: [rerun.io/viewer](https://rerun.io/viewer/pr/6499?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG * [x] If applicable, add a new check to the [release checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)! - [PR Build Summary](https://build.rerun.io/pr/6499) - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html) To run all checks from `main`, comment on the PR with `@rerun-bot full-check`.
rerun-io · Jun 5, 2024 · 4b4cded · 4b4cded
1 parent 7b30934
commit 4b4cded
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 100 deletions.
diff --git a/crates/re_viewer_context/src/test_context.rs b/crates/re_viewer_context/src/test_context.rs
@@ -1,10 +1,11 @@
 use crate::{
-    command_channel, ApplicationSelectionState, ComponentUiRegistry, StoreContext, ViewerContext,
+    command_channel, ApplicationSelectionState, ComponentUiRegistry, RecordingConfig,
+    SpaceViewClassRegistry, StoreContext, ViewerContext,
 };
 
 use re_data_store::LatestAtQuery;
 use re_entity_db::EntityDb;
-use re_log_types::{StoreId, StoreKind};
+use re_log_types::{StoreId, StoreKind, Timeline};
 
 /// Harness to execute code that rely on [`crate::ViewerContext`].
 ///
@@ -19,19 +20,24 @@ use re_log_types::{StoreId, StoreKind};
 /// });
 /// ```
 pub struct TestContext {
-    recording_store: EntityDb,
-    blueprint_store: EntityDb,
-    selection_state: ApplicationSelectionState,
+    pub recording_store: EntityDb,
+    pub blueprint_store: EntityDb,
+    pub space_view_class_registry: SpaceViewClassRegistry,
+    pub selection_state: ApplicationSelectionState,
+    pub active_timeline: Timeline,
 }
 
 impl Default for TestContext {
     fn default() -> Self {
         let recording_store = EntityDb::new(StoreId::random(StoreKind::Recording));
         let blueprint_store = EntityDb::new(StoreId::random(StoreKind::Blueprint));
+        let active_timeline = Timeline::new("time", re_log_types::TimeType::Time);
         Self {
             recording_store,
             blueprint_store,
+            space_view_class_registry: Default::default(),
             selection_state: Default::default(),
+            active_timeline,
         }
     }
 }
@@ -47,10 +53,7 @@ impl TestContext {
     pub fn run(&self, mut func: impl FnMut(&ViewerContext<'_>, &mut egui::Ui)) {
         egui::__run_test_ui(|ui| {
             let re_ui = re_ui::ReUi::load_and_apply(ui.ctx());
-            let blueprint_query = LatestAtQuery::latest(re_log_types::Timeline::new(
-                "timeline",
-                re_log_types::TimeType::Time,
-            ));
+            let blueprint_query = LatestAtQuery::latest(self.active_timeline);
             let (command_sender, _) = command_channel();
             let component_ui_registry = ComponentUiRegistry::new(Box::new(
                 |_ctx, _ui, _ui_layout, _query, _db, _entity_path, _component, _instance| {},
@@ -65,16 +68,19 @@ impl TestContext {
                 hub: &Default::default(),
             };
 
+            let rec_cfg = RecordingConfig::default();
+            rec_cfg.time_ctrl.write().set_timeline(self.active_timeline);
+
             let ctx = ViewerContext {
                 app_options: &Default::default(),
                 cache: &Default::default(),
                 component_ui_registry: &component_ui_registry,
-                space_view_class_registry: &Default::default(),
+                space_view_class_registry: &self.space_view_class_registry,
                 store_context: &store_context,
                 applicable_entities_per_visualizer: &Default::default(),
                 indicated_entities_per_visualizer: &Default::default(),
                 query_results: &Default::default(),
-                rec_cfg: &Default::default(),
+                rec_cfg: &rec_cfg,
                 blueprint_cfg: &Default::default(),
                 selection_state: &self.selection_state,
                 blueprint_query: &blueprint_query,

diff --git a/crates/re_viewport_blueprint/src/space_view.rs b/crates/re_viewport_blueprint/src/space_view.rs
@@ -473,11 +473,13 @@ mod tests {
     };
     use re_types::{archetypes::Points3D, ComponentBatch, ComponentName, Loggable as _};
     use re_viewer_context::{
-        blueprint_timeline, IndicatedEntities, OverridePath, PerVisualizer, SpaceViewClassRegistry,
-        StoreContext, VisualizableEntities,
+        test_context::TestContext, IndicatedEntities, OverridePath, PerVisualizer, StoreContext,
+        VisualizableEntities,
     };
     use std::collections::HashMap;
 
+    use crate::space_view_contents::DataQueryPropertyResolver;
+
     use super::*;
 
     fn save_override(props: EntityProperties, path: &EntityPath, store: &mut EntityDb) {
@@ -495,10 +497,7 @@ mod tests {
 
     #[test]
     fn test_entity_properties() {
-        let space_view_class_registry = SpaceViewClassRegistry::default();
-        let timeline = Timeline::new("time", re_log_types::TimeType::Time);
-        let mut recording = EntityDb::new(StoreId::random(re_log_types::StoreKind::Recording));
-        let mut blueprint = EntityDb::new(StoreId::random(re_log_types::StoreKind::Blueprint));
+        let mut test_ctx = TestContext::default();
         let legacy_auto_properties = EntityPropertyMap::default();
 
         let points = Points3D::new(vec![[1.0, 2.0, 3.0]]);
@@ -510,7 +509,7 @@ mod tests {
         ] {
             let row =
                 DataRow::from_archetype(RowId::new(), TimePoint::default(), path, &points).unwrap();
-            recording.add_data_row(row).ok();
+            test_ctx.recording_store.add_data_row(row).ok();
         }
 
         let recommended = RecommendedSpaceView::new(
@@ -542,35 +541,23 @@ mod tests {
                 .collect(),
         );
 
-        let blueprint_query = LatestAtQuery::latest(blueprint_timeline());
         let contents = &space_view.contents;
 
         let resolver = contents.build_resolver(
-            &space_view_class_registry,
+            &test_ctx.space_view_class_registry,
             &space_view,
             &visualizable_entities,
             &indicated_entities_per_visualizer,
         );
 
         // No overrides set. Everybody has default values.
         {
-            let ctx = StoreContext {
-                app_id: re_log_types::ApplicationId::unknown(),
-                blueprint: &blueprint,
-                default_blueprint: None,
-                recording: &recording,
-                bundle: &Default::default(),
-                hub: &re_viewer_context::StoreHub::test_hub(),
-            };
-
-            let mut query_result = contents.execute_query(&ctx, &visualizable_entities);
-            resolver.update_overrides(
-                &blueprint,
-                &blueprint_query,
-                &timeline,
-                &space_view_class_registry,
+            let query_result = update_overrides(
+                &test_ctx,
+                contents,
+                &visualizable_entities,
+                &resolver,
                 &legacy_auto_properties,
-                &mut query_result,
             );
 
             let parent = query_result
@@ -600,29 +587,18 @@ mod tests {
             save_override(
                 overrides,
                 parent.individual_override_path().unwrap(),
-                &mut blueprint,
+                &mut test_ctx.blueprint_store,
             );
         }
 
         // Parent is not interactive, but children are
         {
-            let ctx = StoreContext {
-                app_id: re_log_types::ApplicationId::unknown(),
-                blueprint: &blueprint,
-                default_blueprint: None,
-                recording: &recording,
-                bundle: &Default::default(),
-                hub: &re_viewer_context::StoreHub::test_hub(),
-            };
-
-            let mut query_result = contents.execute_query(&ctx, &visualizable_entities);
-            resolver.update_overrides(
-                &blueprint,
-                &blueprint_query,
-                &timeline,
-                &space_view_class_registry,
+            let query_result = update_overrides(
+                &test_ctx,
+                contents,
+                &visualizable_entities,
+                &resolver,
                 &legacy_auto_properties,
-                &mut query_result,
             );
 
             let parent_group = query_result
@@ -658,31 +634,19 @@ mod tests {
             save_override(
                 overrides,
                 parent_group.recursive_override_path().unwrap(),
-                &mut blueprint,
+                &mut test_ctx.blueprint_store,
             );
         }
 
         // Nobody is interactive
         {
-            let ctx = StoreContext {
-                app_id: re_log_types::ApplicationId::unknown(),
-                blueprint: &blueprint,
-                default_blueprint: None,
-                recording: &recording,
-                bundle: &Default::default(),
-                hub: &re_viewer_context::StoreHub::test_hub(),
-            };
-
-            let mut query_result = contents.execute_query(&ctx, &visualizable_entities);
-            resolver.update_overrides(
-                &blueprint,
-                &blueprint_query,
-                &timeline,
-                &space_view_class_registry,
+            let query_result = update_overrides(
+                &test_ctx,
+                contents,
+                &visualizable_entities,
+                &resolver,
                 &legacy_auto_properties,
-                &mut query_result,
             );
-
             let parent = query_result
                 .tree
                 .lookup_result_by_path(&EntityPath::from("parent"))
@@ -704,12 +668,9 @@ mod tests {
 
     #[test]
     fn test_component_overrides() {
-        let space_view_class_registry = SpaceViewClassRegistry::default();
-        let timeline = Timeline::new("time", re_log_types::TimeType::Time);
         let legacy_auto_properties = EntityPropertyMap::default();
-        let mut recording = EntityDb::new(StoreId::random(re_log_types::StoreKind::Recording));
-        let mut visualizable_entities_per_visualizer =
-            PerVisualizer::<VisualizableEntities>::default();
+        let mut test_ctx = TestContext::default();
+        let mut visualizable_entities = PerVisualizer::<VisualizableEntities>::default();
 
         // Set up a store DB with some entities.
         {
@@ -726,11 +687,11 @@ mod tests {
                     [&[MyPoint::new(1.0, 2.0)] as _],
                 )
                 .unwrap();
-                recording.add_data_row(row).unwrap();
+                test_ctx.recording_store.add_data_row(row).unwrap();
             }
 
             // All of them are visualizable with some arbitrary visualizer.
-            visualizable_entities_per_visualizer
+            visualizable_entities
                 .0
                 .entry("Points3D".into())
                 .or_insert_with(|| VisualizableEntities(entity_paths.into_iter().collect()));
@@ -750,9 +711,9 @@ mod tests {
         // Things needed to resolve properties:
         let indicated_entities_per_visualizer = PerVisualizer::<IndicatedEntities>::default(); // Don't care about indicated entities.
         let resolver = space_view.contents.build_resolver(
-            &space_view_class_registry,
+            &test_ctx.space_view_class_registry,
             &space_view,
-            &visualizable_entities_per_visualizer,
+            &visualizable_entities,
             &indicated_entities_per_visualizer,
         );
 
@@ -924,7 +885,9 @@ mod tests {
             },
         ) in scenarios.into_iter().enumerate()
         {
-            let mut blueprint = EntityDb::new(StoreId::random(re_log_types::StoreKind::Blueprint));
+            // Reset blueprint store for each scenario.
+            test_ctx.blueprint_store = EntityDb::new(StoreId::random(StoreKind::Blueprint));
+
             let mut add_to_blueprint = |path: &EntityPath, batch: &dyn ComponentBatch| {
                 let row = DataRow::from_component_batches(
                     RowId::new(),
@@ -933,7 +896,7 @@ mod tests {
                     std::iter::once(batch),
                 )
                 .unwrap();
-                blueprint.add_data_row(row).unwrap();
+                test_ctx.blueprint_store.add_data_row(row).unwrap();
             };
 
             // log individual and override components as instructed.
@@ -945,25 +908,12 @@ mod tests {
             }
 
             // Set up a store query and update the overrides.
-            let ctx = StoreContext {
-                app_id: re_log_types::ApplicationId::unknown(),
-                blueprint: &blueprint,
-                default_blueprint: None,
-                recording: &recording,
-                bundle: &Default::default(),
-                hub: &re_viewer_context::StoreHub::test_hub(),
-            };
-            let mut query_result = space_view
-                .contents
-                .execute_query(&ctx, &visualizable_entities_per_visualizer);
-            let blueprint_query = LatestAtQuery::latest(blueprint_timeline());
-            resolver.update_overrides(
-                &blueprint,
-                &blueprint_query,
-                &timeline,
-                &space_view_class_registry,
+            let query_result = update_overrides(
+                &test_ctx,
+                &space_view.contents,
+                &visualizable_entities,
+                &resolver,
                 &legacy_auto_properties,
-                &mut query_result,
             );
 
             // Extract component overrides for testing.
@@ -992,4 +942,36 @@ mod tests {
             assert_eq!(visited, expected_overrides, "Scenario {i}");
         }
     }
+
+    fn update_overrides(
+        test_ctx: &TestContext,
+        contents: &SpaceViewContents,
+        visualizable_entities: &PerVisualizer<VisualizableEntities>,
+        resolver: &DataQueryPropertyResolver<'_>,
+        legacy_auto_properties: &EntityPropertyMap,
+    ) -> re_viewer_context::DataQueryResult {
+        let store_ctx = StoreContext {
+            app_id: re_log_types::ApplicationId::unknown(),
+            blueprint: &test_ctx.blueprint_store,
+            default_blueprint: None,
+            recording: &test_ctx.recording_store,
+            bundle: &Default::default(),
+            hub: &re_viewer_context::StoreHub::test_hub(),
+        };
+
+        let mut query_result = contents.execute_query(&store_ctx, visualizable_entities);
+
+        test_ctx.run(|ctx, _ui| {
+            resolver.update_overrides(
+                ctx.blueprint_db(),
+                ctx.blueprint_query,
+                &test_ctx.active_timeline,
+                ctx.space_view_class_registry,
+                legacy_auto_properties,
+                &mut query_result,
+            );
+        });
+
+        query_result
+    }
 }