1 files changed, 48 insertions, 8 deletions

diff --git a/tests/src/tests.rs b/tests/src/tests.rs
index 46bf343c..da07c628 100644
--- a/tests/src/tests.rs
+++ b/tests/src/tests.rs
@@ -11,6 +11,7 @@ use std::path::{Path, PathBuf};
 use comemo::{Prehashed, Track};
 use elsa::FrozenVec;
 use once_cell::unsync::OnceCell;
+use oxipng::{InFile, Options, OutFile};
 use tiny_skia as sk;
 use typst::diag::{bail, FileError, FileResult};
 use typst::doc::{Document, Frame, FrameItem, Meta};
@@ -77,14 +78,27 @@ fn main() {
         let pdf_path =
             args.pdf.then(|| Path::new(PDF_DIR).join(path).with_extension("pdf"));
 
-        ok += test(&mut world, &src_path, &png_path, &ref_path, pdf_path.as_deref())
-            as usize;
+        ok += test(
+            &mut world,
+            &src_path,
+            &png_path,
+            &ref_path,
+            pdf_path.as_deref(),
+            args.update,
+        ) as usize;
     }
 
     if len > 1 {
         println!("{ok} / {len} tests passed.");
     }
 
+    if ok != len {
+        println!(
+            "Set the UPDATE_EXPECT environment variable or pass the \
+             --update flag to update the reference image(s)."
+        );
+    }
+
     if ok < len {
         std::process::exit(1);
     }
@@ -95,6 +109,7 @@ struct Args {
     filter: Vec<String>,
     exact: bool,
     pdf: bool,
+    update: bool,
     print: PrintConfig,
 }
 
@@ -111,6 +126,7 @@ impl Args {
         let mut filter = Vec::new();
         let mut exact = false;
         let mut pdf = false;
+        let mut update = env::var_os("UPDATE_EXPECT").is_some();
         let mut print = PrintConfig::default();
 
         for arg in args {
@@ -121,6 +137,8 @@ impl Args {
                 "--exact" => exact = true,
                 // Generate PDFs.
                 "--pdf" => pdf = true,
+                // Update the reference images.
+                "--update" => update = true,
                 // Debug print the syntax trees.
                 "--syntax" => print.syntax = true,
                 // Debug print the model.
@@ -132,7 +150,7 @@ impl Args {
             }
         }
 
-        Self { filter, exact, pdf, print }
+        Self { filter, exact, pdf, update, print }
     }
 
     fn matches(&self, path: &Path) -> bool {
@@ -339,6 +357,7 @@ fn test(
     png_path: &Path,
     ref_path: &Path,
     pdf_path: Option<&Path>,
+    update: bool,
 ) -> bool {
     let name = src_path.strip_prefix(TYP_DIR).unwrap_or(src_path);
     println!("Testing {}", name.display());
@@ -346,6 +365,7 @@ fn test(
     let text = fs::read_to_string(src_path).unwrap();
 
     let mut ok = true;
+    let mut updated = false;
     let mut frames = vec![];
     let mut line = 0;
     let mut compare_ref = true;
@@ -404,16 +424,26 @@ fn test(
                     .zip(ref_pixmap.data())
                     .any(|(&a, &b)| a.abs_diff(b) > 2)
             {
-                println!("  Does not match reference image. ❌");
-                ok = false;
+                if update {
+                    update_image(png_path, ref_path);
+                    updated = true;
+                } else {
+                    println!("  Does not match reference image. ❌");
+                    ok = false;
+                }
             }
         } else if !document.pages.is_empty() {
-            println!("  Failed to open reference image. ❌");
-            ok = false;
+            if update {
+                update_image(png_path, ref_path);
+                updated = true;
+            } else {
+                println!("  Failed to open reference image. ❌");
+                ok = false;
+            }
         }
     }
 
-    if ok {
+    if ok && !updated {
         if world.print == PrintConfig::default() {
             print!("\x1b[1A");
         }
@@ -423,6 +453,16 @@ fn test(
     ok
 }
 
+fn update_image(png_path: &Path, ref_path: &Path) {
+    println!("  Updated reference image. ✔");
+    oxipng::optimize(
+        &InFile::Path(png_path.to_owned()),
+        &OutFile::Path(Some(ref_path.to_owned())),
+        &Options::max_compression(),
+    )
+    .unwrap();
+}
+
 fn test_part(
     world: &mut TestWorld,
     src_path: &Path,