Add screenshot capabilities

CHANGELOG.md

@@ -9,6 +9,9 @@
 
 ## Unreleased
 
+* Add screenshot instructions to Toolproof
+* Add `extract` concept to pull retrievals to disk
+
 ## v0.6.1 (November 28, 2024)
 
 * Log inner macro steps

docs/content/docs/_index.md

@@ -135,6 +135,20 @@ steps:
 In future runs, Toolproof will ensure the retrieved value matches the `snapshot_content` key. Running Toolproof in
 interactive mode (`-i`) will also allow you to accept the changes and update the file automatically.
 
+### Snapshots
+
+Any Retrieval can also drive an extract. To do so, place the step inside an object under a `extract` key, alongside an `extract_location` key:
+```yaml
+steps:
+  - extract: stdout
+    extract_location: "%toolproof_process_directory%/extracted_stdout.txt"
+```
+
+After running Toolproof, the value will be written to that file.
+
+Toolproof never reads this file, so this step doesn't have any bearing on the success of the test.
+Instead, this is intended to pull information from tests to use in other tooling.
+
 ## Test environment
 
 Toolproof automatically runs tests in a temporary directory that is discarded at the end of a run.

docs/content/docs/functions.md

@@ -7,7 +7,7 @@ weight: 6
 
 Toolproof provides the following Instructions:
 
-## Filesystem 
+## Filesystem
 
 Instructions:
 - `I have a {filename} file with the content {contents}`
@@ -16,7 +16,7 @@ Retrievals:
 - `The file {filename}`
   - Returns a string value
 
-## Process 
+## Process
 
 Instructions:
 - `I have the environment variable {name} set to {value}`
@@ -29,16 +29,18 @@ Retrievals:
 - `stderr`
   - Returns a string value
 
-## Hosting 
+## Hosting
 
 Instructions:
 - `I serve the directory {dir}`
 
-## Browser 
+## Browser
 
 Instructions:
 - `In my browser, I load {url}`
 - `In my browser, I evaluate {js}`
+- `In my browser, I screenshot the viewport to {filepath}`
+- `In my browser, I screenshot the element {selector} to {filepath}`
 
 Retrievals:
 - `In my browser, the result of {js}`

toolproof/src/civilization.rs

@@ -74,6 +74,10 @@ impl<'u> Civilization<'u> {
         tmp_dir.join(PathBuf::from(filename))
     }
 
+    pub fn ensure_path(&mut self, file_path: &PathBuf) {
+        fs::create_dir_all(file_path.parent().unwrap()).unwrap();
+    }
+
     pub fn write_file(&mut self, filename: &str, contents: &str) {
         let file_path = self.tmp_file_path(filename);
         fs::create_dir_all(file_path.parent().unwrap()).unwrap();

toolproof/src/definitions/browser/mod.rs

@@ -1,13 +1,21 @@
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use chromiumoxide::cdp::browser_protocol::page::{
+    CaptureScreenshotFormat, CaptureScreenshotParams,
+};
 use chromiumoxide::cdp::browser_protocol::target::CreateTargetParams;
 use chromiumoxide::error::CdpError;
+use chromiumoxide::handler::viewport::Viewport;
+use chromiumoxide::page::ScreenshotParams;
 use futures::StreamExt;
 use tokio::task::JoinHandle;
 
 use crate::civilization::Civilization;
-use crate::errors::{ToolproofInputError, ToolproofInternalError, ToolproofStepError};
+use crate::errors::{
+    ToolproofInputError, ToolproofInternalError, ToolproofStepError, ToolproofTestFailure,
+};
 use crate::options::ToolproofParams;
 
 use super::{SegmentArgs, ToolproofInstruction, ToolproofRetriever};
@@ -34,7 +42,21 @@ async fn try_launch_browser(mut max: usize) -> (Browser, chromiumoxide::Handler)
     let mut launch = Err(CdpError::NotFound);
     while launch.is_err() && max > 0 {
         max -= 1;
-        launch = Browser::launch(BrowserConfig::builder().build().unwrap()).await;
+        launch = Browser::launch(
+            BrowserConfig::builder()
+                .headless_mode(chromiumoxide::browser::HeadlessMode::New)
+                .viewport(Some(Viewport {
+                    width: 1600,
+                    height: 900,
+                    device_scale_factor: Some(2.0),
+                    emulating_mobile: false,
+                    is_landscape: true,
+                    has_touch: false,
+                }))
+                .build()
+                .unwrap(),
+        )
+        .await;
     }
     match launch {
         Ok(res) => res,
@@ -44,6 +66,29 @@ async fn try_launch_browser(mut max: usize) -> (Browser, chromiumoxide::Handler)
     }
 }
 
+fn chrome_image_format(filepath: &PathBuf) -> Result<CaptureScreenshotFormat, ToolproofStepError> {
+    match filepath.extension() {
+        Some(ext) => {
+            let ext = ext.to_string_lossy().to_lowercase();
+            match ext.as_str() {
+                "png" => Ok(CaptureScreenshotFormat::Png),
+                "webp" => Ok(CaptureScreenshotFormat::Webp),
+                "jpg" | "jpeg" => Ok(CaptureScreenshotFormat::Jpeg),
+                _ => Err(ToolproofStepError::External(
+                    ToolproofInputError::StepRequirementsNotMet {
+                        reason: "Image file extension must be png, webp, jpeg, or jpg".to_string(),
+                    },
+                )),
+            }
+        }
+        None => Err(ToolproofStepError::External(
+            ToolproofInputError::StepRequirementsNotMet {
+                reason: "Image file path must have an extension".to_string(),
+            },
+        )),
+    }
+}
+
 impl BrowserTester {
     async fn initialize(params: &ToolproofParams) -> Self {
         match params.browser {
@@ -145,6 +190,63 @@ impl BrowserWindow {
                 .map_err(|inner| ToolproofStepError::Internal(inner.into())),
         }
     }
+
+    async fn screenshot_page(&self, filepath: PathBuf) -> Result<(), ToolproofStepError> {
+        match self {
+            BrowserWindow::Chrome(page) => {
+                let image_format = chrome_image_format(&filepath)?;
+
+                page.save_screenshot(
+                    ScreenshotParams {
+                        cdp_params: CaptureScreenshotParams {
+                            format: Some(image_format),
+                            ..CaptureScreenshotParams::default()
+                        },
+                        full_page: Some(false),
+                        omit_background: Some(false),
+                    },
+                    filepath,
+                )
+                .await
+                .map(|_| ())
+                .map_err(|e| ToolproofStepError::Internal(e.into()))
+            }
+            BrowserWindow::Pagebrowse(_) => Err(ToolproofStepError::Internal(
+                ToolproofInternalError::Custom {
+                    msg: "Screenshots not yet implemented for Pagebrowse".to_string(),
+                },
+            )),
+        }
+    }
+
+    async fn screenshot_element(
+        &self,
+        selector: &str,
+        filepath: PathBuf,
+    ) -> Result<(), ToolproofStepError> {
+        match self {
+            BrowserWindow::Chrome(page) => {
+                let image_format = chrome_image_format(&filepath)?;
+
+                let element = page.find_element(selector).await.map_err(|e| {
+                    ToolproofStepError::Assertion(ToolproofTestFailure::Custom {
+                        msg: format!("Element {selector} could not be screenshot: {e}"),
+                    })
+                })?;
+
+                element
+                    .save_screenshot(image_format, filepath)
+                    .await
+                    .map(|_| ())
+                    .map_err(|e| ToolproofStepError::Internal(e.into()))
+            }
+            BrowserWindow::Pagebrowse(_) => Err(ToolproofStepError::Internal(
+                ToolproofInternalError::Custom {
+                    msg: "Screenshots not yet implemented for Pagebrowse".to_string(),
+                },
+            )),
+        }
+    }
 }
 
 mod load_page {
@@ -318,3 +420,76 @@ mod eval_js {
         }
     }
 }
+
+mod screenshots {
+    use crate::errors::{ToolproofInternalError, ToolproofTestFailure};
+
+    use super::*;
+
+    pub struct ScreenshotViewport;
+
+    inventory::submit! {
+        &ScreenshotViewport as &dyn ToolproofInstruction
+    }
+
+    #[async_trait]
+    impl ToolproofInstruction for ScreenshotViewport {
+        fn segments(&self) -> &'static str {
+            "In my browser, I screenshot the viewport to {filepath}"
+        }
+
+        async fn run(
+            &self,
+            args: &SegmentArgs<'_>,
+            civ: &mut Civilization,
+        ) -> Result<(), ToolproofStepError> {
+            let filepath = args.get_string("filepath")?;
+            let resolved_path = civ.tmp_file_path(&filepath);
+            civ.ensure_path(&resolved_path);
+
+            let Some(window) = civ.window.as_ref() else {
+                return Err(ToolproofStepError::External(
+                    ToolproofInputError::StepRequirementsNotMet {
+                        reason: "no page has been loaded into the browser for this test".into(),
+                    },
+                ));
+            };
+
+            window.screenshot_page(resolved_path).await
+        }
+    }
+
+    pub struct ScreenshotElement;
+
+    inventory::submit! {
+        &ScreenshotElement as &dyn ToolproofInstruction
+    }
+
+    #[async_trait]
+    impl ToolproofInstruction for ScreenshotElement {
+        fn segments(&self) -> &'static str {
+            "In my browser, I screenshot the element {selector} to {filepath}"
+        }
+
+        async fn run(
+            &self,
+            args: &SegmentArgs<'_>,
+            civ: &mut Civilization,
+        ) -> Result<(), ToolproofStepError> {
+            let selector = args.get_string("selector")?;
+            let filepath = args.get_string("filepath")?;
+            let resolved_path = civ.tmp_file_path(&filepath);
+            civ.ensure_path(&resolved_path);
+
+            let Some(window) = civ.window.as_ref() else {
+                return Err(ToolproofStepError::External(
+                    ToolproofInputError::StepRequirementsNotMet {
+                        reason: "no page has been loaded into the browser for this test".into(),
+                    },
+                ));
+            };
+
+            window.screenshot_element(&selector, resolved_path).await
+        }
+    }
+}