test: Refactor book snapshot tests

book/tests/snapshot.rs

@@ -1,6 +1,6 @@
 #![cfg(not(target_family = "wasm"))]
 /// This test:
-/// - Extracts PRQL code blocks into the `examples` path.
+/// - Extracts PRQL code blocks into files in the `examples` path
 /// - Converts them to SQL using insta, raising an error if there's a diff.
 /// - Replaces the PRQL code block with a comparison table.
 ///
@@ -18,69 +18,38 @@
 // us. They introduce a bunch of non-rust dependencies, which is not ideal, but
 // passable. They don't let us customize our formatting (e.g. in a table).
 //
-use anyhow::{bail, Result};
+use anyhow::{bail, Error, Result};
 use globset::Glob;
-use insta::{assert_display_snapshot, assert_snapshot, glob};
+use insta::{assert_snapshot, glob};
+use itertools::Itertools;
 use log::warn;
 use prql_compiler::*;
 use pulldown_cmark::{CodeBlockKind, Event, Parser, Tag};
-use std::fs;
 use std::path::Path;
+use std::{collections::HashMap, fs};
 use walkdir::WalkDir;
 
 #[test]
-fn run_examples() -> Result<()> {
-    // TODO: In CI this could pass by replacing incorrect files. To catch that,
-    // we could check if there are any diffs after this has run?
-
+fn test_examples() -> Result<()> {
     // Note that on windows, markdown is read differently, and so
     // writing on Windows. ref https://github.com/PRQL/prql/issues/356
     #[cfg(not(target_family = "windows"))]
-    write_reference_prql()?;
-    run_reference_prql();
-
-    // TODO: Currently we run this in the same test, since we need the
-    // `write_reference_prql` function to have been run. If we could iterate
-    // over the PRQL examples without writing them to disk, we could run this as
-    // a separate test. (Though then we'd lose the deferred failures feature
-    // that insta's `glob!` macro provides.)
-    run_display_reference_prql();
+    write_prql_examples()?;
+    test_prql_examples();
 
     Ok(())
 }
 
-/// Extract reference examples from the PRQL docs and write them to the
-/// `tests/prql` path, one in each file.
-// We could alternatively have used something like
-// https://github.com/earldouglas/codedown, but it's not much code, and it
-// requires no dependencies.
-//
-// We allow dead_code because of the window issue described above. (Can we allow
-// it only for windows?)
-#[allow(dead_code)]
-fn write_reference_prql() -> Result<()> {
-    // Remove old .prql files, since we're going to rewrite them, and we don't want
-    // old files which wouldn't be rewritten from hanging around.
-    // We use `trash`, since we don't want to be removing files with test code
-    // in case there's a bug.
-    //
-    // A more elegant approach would be to keep a list of the files and remove
-    // the ones we don't write.
-
-    let examples_path = Path::new("tests/prql");
-    if examples_path.exists() {
-        trash::delete(Path::new("tests/prql")).unwrap_or_else(|e| {
-            warn!("Failed to delete old examples: {}", e);
-        });
-    }
+const ROOT_EXAMPLES_PATH: &str = "tests/prql";
 
+/// Collect all the PRQL examples in the book, as a map of <Path, PRQL>.
+fn collect_book_examples() -> Result<HashMap<String, String>> {
     let glob = Glob::new("**/*.md")?.compile_matcher();
-
-    WalkDir::new(Path::new("./src/"))
+    let examples_in_book: HashMap<String, String> = WalkDir::new(Path::new("./src/"))
         .into_iter()
         .flatten()
         .filter(|x| glob.is_match(x.path()))
-        .try_for_each(|dir_entry| {
+        .flat_map(|dir_entry| {
             let text = fs::read_to_string(dir_entry.path())?;
             let mut parser = Parser::new(&text);
             let mut prql_blocks = vec![];
@@ -102,33 +71,85 @@ fn write_reference_prql() -> Result<()> {
                     _ => {}
                 }
             }
-
-            // Write each one to a new file.
-            prql_blocks
+            let snapshot_prefix = &dir_entry
+                .path()
+                .strip_prefix("./src/")?
+                .to_str()
+                .unwrap()
+                .trim_end_matches(".md");
+            Ok(prql_blocks
                 .iter()
                 .enumerate()
-                .try_for_each(|(i, example)| {
-                    let file_relative = &dir_entry
-                        .path()
-                        .strip_prefix("./src/")?
-                        .to_str()
-                        .unwrap()
-                        .trim_end_matches(".md");
-                    let prql_path = format!("tests/prql/{file_relative}-{i}.prql");
-
-                    fs::create_dir_all(Path::new(&prql_path).parent().unwrap())?;
-                    fs::write(prql_path, example.to_string())?;
-
-                    Ok::<(), anyhow::Error>(())
-                })?;
-            Ok(())
+                .map(|(i, example)| {
+                    (
+                        format!("{ROOT_EXAMPLES_PATH}/{snapshot_prefix}-{i}.prql"),
+                        example.to_string(),
+                    )
+                })
+                .collect::<HashMap<_, _>>())
+        })
+        .flatten()
+        .collect();
+
+    Ok(examples_in_book)
+}
+
+/// Extract reference examples from the PRQL docs and write them to the
+/// `tests/prql` path, one in each file.
+// We could alternatively have used something like
+// https://github.com/earldouglas/codedown, but it's not much code, and it
+// requires no dependencies.
+//
+// We allow dead_code because of the window issue described above. (Can we allow
+// it only for windows?)
+#[allow(dead_code)]
+fn write_prql_examples() -> Result<()> {
+    // If we have to modify any files, raise an error at the end, so it fails in CI.
+    let mut snapshots_updated = false;
+
+    let examples_path = Path::new(ROOT_EXAMPLES_PATH);
+    let prql_matcher = Glob::new("**.prql")?.compile_matcher();
+    let mut existing_snapshots: HashMap<_, _> = WalkDir::new(examples_path)
+        .into_iter()
+        .flatten()
+        .filter(|x| prql_matcher.is_match(x.path()))
+        .map(|x| x.path().to_owned())
+        .map(|x| Ok::<_, Error>((x.to_string_lossy().to_string(), fs::read_to_string(x)?)))
+        .try_collect()?;
+
+    // Write any new snapshots, or update any that have changed. =
+    collect_book_examples()?
+        .iter()
+        .try_for_each(|(prql_path, example)| {
+            if existing_snapshots
+                .remove(prql_path)
+                .map(|existing| existing != *example)
+                .unwrap_or(true)
+            {
+                snapshots_updated = true;
+                fs::create_dir_all(Path::new(prql_path).parent().unwrap())?;
+                fs::write(prql_path, example)?;
+            }
+
+            Ok::<(), anyhow::Error>(())
         })?;
 
+    // If there are any files left in `existing_snapshots`, we remove them, since
+    // they don't reference anything.
+    existing_snapshots.iter().for_each(|(path, _)| {
+        trash::delete(path).unwrap_or_else(|e| {
+            warn!("Failed to delete unreferenced example: {}", e);
+        })
+    });
+
+    if snapshots_updated {
+        bail!("Some book snapshots were not consistent with the queries in the book and have been updated. Subsequent runs should pass.");
+    }
     Ok(())
 }
 
 /// Snapshot the output of each example.
-fn run_reference_prql() {
+fn test_prql_examples() {
     glob!("prql/**/*.prql", |path| {
         let prql = fs::read_to_string(path).unwrap();
 
@@ -147,20 +168,25 @@ fn run_reference_prql() {
 }
 
 /// Snapshot the display trait output of each example.
-// Currently not a separate test, see notes in caller.
 //
 // TODO: this involves writing out almost the same PRQL again — instead we could
 // compare the output of Display to the auto-formatted source. But we need an
 // autoformatter for that (unless we want to raise on any non-matching input,
 // which seems very strict)
-fn run_display_reference_prql() {
-    glob!("prql/**/*.prql", |path| {
-        let prql = fs::read_to_string(path).unwrap();
-
-        if prql.contains("skip_test") {
-            return;
-        }
+#[test]
+fn test_display() -> Result<(), ErrorMessages> {
+    use prql_compiler::downcast;
+    collect_book_examples()
+        .map_err(downcast)?
+        .iter()
+        .try_for_each(|(path, example)| {
+            assert_snapshot!(
+                path.to_owned(),
+                prql_to_pl(example).and_then(pl_to_prql)?,
+                example
+            );
+            Ok::<(), ErrorMessages>(())
+        })?;
 
-        assert_display_snapshot!(prql_to_pl(&prql).and_then(pl_to_prql).unwrap());
-    });
+    Ok(())
 }