Handle ViewportCommand::Screenshot in kittest (#8256)

Makes it possible to use egui_mcp to drive apps headlessly via kittest

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"] }
 

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,

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<Arc<ColorImage>>,
+    }
+
+    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());
+}