From 600286d05606cfaecd577f74eb2cb7edeac234bd Mon Sep 17 00:00:00 2001 From: Lucas Meurer Date: Wed, 24 Jun 2026 19:58:30 +0200 Subject: [PATCH] Handle `ViewportCommand::Screenshot` in kittest (#8256) Makes it possible to use egui_mcp to drive apps headlessly via kittest --- Cargo.lock | 1 + crates/egui_kittest/Cargo.toml | 1 + crates/egui_kittest/src/lib.rs | 70 +++++++++++++++++++++++++ crates/egui_kittest/tests/screenshot.rs | 63 ++++++++++++++++++++++ 4 files changed, 135 insertions(+) create mode 100644 crates/egui_kittest/tests/screenshot.rs diff --git a/Cargo.lock b/Cargo.lock index 30b9ad9f0..1ae343af4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1468,6 +1468,7 @@ dependencies = [ "egui_extras", "image", "kittest", + "log", "open", "pollster", "serde", diff --git a/crates/egui_kittest/Cargo.toml b/crates/egui_kittest/Cargo.toml index 2d26e6bd8..a770438c3 100644 --- a/crates/egui_kittest/Cargo.toml +++ b/crates/egui_kittest/Cargo.toml @@ -37,6 +37,7 @@ x11 = ["eframe?/x11"] egui.workspace = true eframe = { workspace = true, optional = true } kittest.workspace = true +log.workspace = true serde.workspace = true toml = { workspace = true, features = ["parse", "serde"] } diff --git a/crates/egui_kittest/src/lib.rs b/crates/egui_kittest/src/lib.rs index 4192a103c..3c859560d 100644 --- a/crates/egui_kittest/src/lib.rs +++ b/crates/egui_kittest/src/lib.rs @@ -175,6 +175,11 @@ impl<'a, State> Harness<'a, State> { #[cfg(feature = "snapshot")] snapshot_results: SnapshotResults::default(), }; + // Fulfill any screenshot requested during the initial frame above (which didn't go + // through `_step`). + #[cfg(any(feature = "wgpu", feature = "snapshot"))] + harness.handle_screenshots(); + // Run the harness until it is stable, ensuring that all Areas are shown and animations are done harness.run_ok(); harness @@ -274,6 +279,9 @@ impl<'a, State> Harness<'a, State> { ); self.renderer.handle_delta(&output.textures_delta); self.output = output; + + #[cfg(any(feature = "wgpu", feature = "snapshot"))] + self.handle_screenshots(); } /// Calculate the rect that includes all popups and tooltips. @@ -667,6 +675,56 @@ impl<'a, State> Harness<'a, State> { self.renderer.render(&self.ctx, &output) } + /// Fulfill any [`egui::ViewportCommand::Screenshot`] requests made by the app during the + /// last frame. + /// + /// If a screenshot was requested and no renderer is available, an error will be logged. + #[cfg(any(feature = "wgpu", feature = "snapshot"))] + fn handle_screenshots(&mut self) { + // Collect all screenshot requests from this frame's viewport output. + let requests: Vec<(ViewportId, egui::UserData)> = self + .output + .viewport_output + .iter() + .flat_map(|(id, viewport)| { + viewport.commands.iter().filter_map(move |command| { + if let egui::ViewportCommand::Screenshot(user_data) = command { + Some((*id, user_data.clone())) + } else { + None + } + }) + }) + .collect(); + + if requests.is_empty() { + return; + } + + // Render the frame once and reuse it for every request. We render without the synthetic + // mouse cursor since a real screenshot wouldn't include the OS cursor either. + let image = match self.renderer.render(&self.ctx, &self.output) { + Ok(image) => image, + Err(err) => { + log::error!("Failed to render screenshot requested via ViewportCommand: {err}"); + return; + } + }; + let image = std::sync::Arc::new(rgba_image_to_color_image(&image)); + + for (viewport_id, user_data) in requests { + self.input.events.push(egui::Event::Screenshot { + viewport_id, + user_data, + image: std::sync::Arc::clone(&image), + }); + } + + // Make sure the run loop runs at least one more frame so the app actually receives the + // queued screenshot event. + self.ctx.request_repaint(); + } + /// Get the root viewport output fn root_viewport_output(&self) -> &egui::ViewportOutput { self.output @@ -810,6 +868,18 @@ impl<'a> Harness<'a> { } } +/// Convert a rendered [`image::RgbaImage`] (premultiplied alpha, as produced by the renderer) +/// into an [`egui::ColorImage`] suitable for [`egui::Event::Screenshot`]. +#[cfg(any(feature = "wgpu", feature = "snapshot"))] +fn rgba_image_to_color_image(image: &image::RgbaImage) -> egui::ColorImage { + let size = [image.width() as usize, image.height() as usize]; + let pixels = image + .pixels() + .map(|p| Color32::from_rgba_unmultiplied(p[0], p[1], p[2], p[3])) + .collect(); + egui::ColorImage::new(size, pixels) +} + impl<'tree, 'node, State> Queryable<'tree, 'node, Node<'tree>> for Harness<'_, State> where 'node: 'tree, diff --git a/crates/egui_kittest/tests/screenshot.rs b/crates/egui_kittest/tests/screenshot.rs new file mode 100644 index 000000000..07c971577 --- /dev/null +++ b/crates/egui_kittest/tests/screenshot.rs @@ -0,0 +1,63 @@ +#![cfg(feature = "wgpu")] + +use std::sync::Arc; + +use egui::{Color32, ColorImage, Vec2}; +use egui_kittest::Harness; + +/// Requesting a screenshot via [`egui::ViewportCommand::Screenshot`] from within the app should +/// be fulfilled by the harness (when rendering is enabled) and delivered back via +/// [`egui::Event::Screenshot`]. +#[test] +#[cfg(any(feature = "wgpu", feature = "snapshot"))] +fn screenshot_viewport_command() { + #[derive(Default)] + struct State { + requested: bool, + screenshot: Option>, + } + + let mut harness = Harness::builder() + .with_size(Vec2::new(100.0, 80.0)) + .build_ui_state( + |ui, state: &mut State| { + // Paint the whole content area with a known color so we can verify the capture. + ui.painter() + .rect_filled(ui.ctx().content_rect(), 0.0, Color32::RED); + + // Request a screenshot once. + if !state.requested { + state.requested = true; + ui.ctx() + .send_viewport_cmd(egui::ViewportCommand::Screenshot(Default::default())); + } + + // Capture the screenshot once it's delivered. + ui.input(|i| { + for event in &i.raw.events { + if let egui::Event::Screenshot { image, .. } = event { + state.screenshot = Some(Arc::clone(image)); + } + } + }); + }, + State::default(), + ); + + harness.run(); + + let screenshot = harness + .state() + .screenshot + .clone() + .expect("Expected a screenshot to be delivered via Event::Screenshot"); + + // The frame was filled with red, so the center pixel should be red. + let center = screenshot.pixels[screenshot.pixels.len() / 2]; + assert_eq!(center, Color32::RED, "center pixel should be red"); + + // The screenshot should match the rendered frame size. + let rendered = harness.render().unwrap(); + assert_eq!(screenshot.width() as u32, rendered.width()); + assert_eq!(screenshot.height() as u32, rendered.height()); +}