re_test_viewport/
lib.rs

1//! Extends the `re_test_context` with viewport-related features.
2
3mod test_view;
4
5use ahash::HashMap;
6use re_test_context::TestContext;
7use re_test_context::external::egui_kittest::{SnapshotOptions, SnapshotResult};
8use re_viewer_context::{Contents, ViewId, ViewerContext, VisitorControlFlow};
9use re_viewport::execute_systems_for_view;
10use re_viewport_blueprint::{ViewBlueprint, ViewportBlueprint};
11pub use test_view::TestView;
12
13/// Extension trait to [`TestContext`] for blueprint-related features.
14pub trait TestContextExt {
15    /// See docstring on the implementation below.
16    fn setup_viewport_blueprint<R>(
17        &mut self,
18        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R,
19    ) -> R;
20
21    /// Displays the UI for a single given view.
22    fn ui_for_single_view(&self, ui: &mut egui::Ui, ctx: &ViewerContext<'_>, view_id: ViewId);
23
24    /// [`TestContext::run`] inside a central panel that displays the ui for a single given view.
25    fn run_with_single_view(&self, ui: &mut egui::Ui, view_id: ViewId);
26
27    fn run_view_ui_and_save_snapshot(
28        &self,
29        view_id: ViewId,
30        snapshot_name: &str,
31        size: egui::Vec2,
32        snapshot_options: Option<SnapshotOptions>,
33    ) -> SnapshotResult;
34}
35
36impl TestContextExt for TestContext {
37    /// Inspect or update the blueprint of a [`TestContext`].
38    ///
39    /// This helper works by deserializing the current blueprint, providing it to the provided
40    /// closure, and saving it back to the blueprint store. The closure should call the appropriate
41    /// methods of [`ViewportBlueprint`] to inspect and/or create views and containers as required.
42    ///
43    /// Each time [`TestContextExt::setup_viewport_blueprint`] is called, it entirely recomputes the
44    /// "query results", i.e., the [`re_viewer_context::DataResult`]s that each view contains, based
45    /// on the current content of the recording store.
46    ///
47    /// Important pre-requisite:
48    /// - The current timeline must already be set to the timeline of interest, because some
49    ///   updates are timeline-dependant (in particular those related to visible time range).
50    /// - The view classes used by view must be already registered (see
51    ///   [`TestContext::register_view_class`]).
52    /// - The data store must be already populated for the views to have any content (see, e.g.,
53    ///   [`TestContext::log_entity`]).
54    ///
55    fn setup_viewport_blueprint<R>(
56        &mut self,
57        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R,
58    ) -> R {
59        let mut setup_blueprint: Option<_> = Some(setup_blueprint);
60
61        let mut result = None;
62
63        egui::__run_test_ctx(|egui_ctx| {
64            // We use `take` to ensure that the blueprint is setup only once, since egui forces
65            // us to a `FnMut` closure.
66            if let Some(setup_blueprint) = setup_blueprint.take() {
67                self.run(egui_ctx, |ctx| {
68                    let mut viewport_blueprint =
69                        ViewportBlueprint::from_db(ctx.blueprint_db(), &self.blueprint_query);
70                    result = Some(setup_blueprint(ctx, &mut viewport_blueprint));
71                    viewport_blueprint.save_to_blueprint_store(ctx);
72                });
73
74                self.handle_system_commands(egui_ctx);
75
76                // Reload the blueprint store and execute all view queries.
77                let blueprint_query = self.blueprint_query.clone();
78                let viewport_blueprint =
79                    ViewportBlueprint::from_db(self.active_blueprint(), &blueprint_query);
80
81                let mut query_results = HashMap::default();
82
83                self.run(egui_ctx, |ctx| {
84                    let _ignored = viewport_blueprint.visit_contents::<()>(&mut |contents, _| {
85                        if let Contents::View(view_id) = contents {
86                            let view_blueprint = viewport_blueprint
87                                .view(view_id)
88                                .expect("view is known to exist");
89
90                            let class_registry = ctx.view_class_registry();
91                            let class_identifier = view_blueprint.class_identifier();
92                            let class = class_registry.class(class_identifier).unwrap_or_else(|| panic!("The class '{class_identifier}' must be registered beforehand"));
93
94                            let visualizable_entities_for_view = ctx.collect_visualizable_entities_for_view_class(class_identifier);
95
96                            let query_range = view_blueprint.query_range(
97                                ctx.blueprint_db(),
98                                ctx.blueprint_query,
99                                ctx.time_ctrl.timeline(),
100                                class_registry,
101                                self.view_states.lock().get_mut_or_create(*view_id, class),
102                            );
103
104                            let data_query_result = view_blueprint.contents.build_data_result_tree(
105                                ctx.store_context,
106                                ctx.time_ctrl.timeline(),
107                                class_registry,
108                                ctx.blueprint_query,
109                                &query_range,
110                                &visualizable_entities_for_view,
111                                ctx.indicated_entities_per_visualizer,
112                                ctx.app_options(),
113                            );
114
115                            query_results.insert(*view_id, data_query_result);
116                        }
117
118                        VisitorControlFlow::Continue
119                    });
120                });
121
122                self.query_results = query_results;
123            }
124        });
125
126        result.expect("The `setup_closure` is expected to be called at least once")
127    }
128
129    /// Displays the UI for a single given view.
130    fn ui_for_single_view(&self, ui: &mut egui::Ui, ctx: &ViewerContext<'_>, view_id: ViewId) {
131        let view_blueprint =
132            ViewBlueprint::try_from_db(view_id, ctx.store_context.blueprint, ctx.blueprint_query)
133                .expect("expected the view id to be known to the blueprint store");
134
135        let class_registry = ctx.view_class_registry();
136        let class_identifier = view_blueprint.class_identifier();
137        let view_class = class_registry.get_class_or_log_error(class_identifier);
138
139        let mut view_states = self.view_states.lock();
140        view_states.reset_visualizer_errors();
141        let view_state = view_states.get_mut_or_create(view_id, view_class);
142
143        let context_system_once_per_frame_results = class_registry
144            .run_once_per_frame_context_systems(ctx, std::iter::once(class_identifier));
145        let (view_query, system_execution_output) = execute_systems_for_view(
146            ctx,
147            &view_blueprint,
148            view_state,
149            &context_system_once_per_frame_results,
150        );
151        view_states.report_visualizer_errors(view_id, &system_execution_output);
152
153        let view_state = view_states.get_mut_or_create(view_id, view_class);
154        view_class
155            .ui(ctx, ui, view_state, &view_query, system_execution_output)
156            .expect("failed to run view ui");
157    }
158
159    /// [`TestContext::run`] for a single view.
160    fn run_with_single_view(&self, ui: &mut egui::Ui, view_id: ViewId) {
161        self.run_ui(ui, |ctx, ui| {
162            self.ui_for_single_view(ui, ctx, view_id);
163        });
164
165        self.handle_system_commands(ui.ctx());
166    }
167
168    fn run_view_ui_and_save_snapshot(
169        &self,
170        view_id: ViewId,
171        snapshot_name: &str,
172        size: egui::Vec2,
173        snapshot_options: Option<SnapshotOptions>,
174    ) -> SnapshotResult {
175        let mut harness = self.setup_kittest_for_rendering_3d(size).build_ui(|ui| {
176            self.run_with_single_view(ui, view_id);
177        });
178        harness.run();
179
180        if let Some(snapshot_options) = snapshot_options {
181            harness.try_snapshot_options(snapshot_name, &snapshot_options)
182        } else {
183            harness.try_snapshot(snapshot_name)
184        }
185    }
186}
re_test_viewport/lib.rs

re_test_viewport/
lib.rs
