Merge pull request #1081 from EnergySystemsModellingLab/regenerate-patch-test-data

Allow for regenerating test data for patched examples + tidy-ups

Author: alexdewar

docs/developer_guide/architecture_quickstart.md

@@ -108,6 +108,22 @@ just regenerate_test_data
 If you do so, please verify that the changes to the output files are at least roughly what was
 expected, before you commit these updated test files.
 
+Note that if you only need to regenerate the data for some of the models, you can specify this with
+additional arguments, e.g.:
+
+```sh
+just regenerate_test_data simple muse1_default
+```
+
+This avoids regenerating data for other models unnecessarily, which can result in negligible
+differences in floating-point values in the output files.
+
+If the model is a [patched example], then you need to pass the `--patch` flag, e.g.:
+
+```sh
+just regenerate_test_data --patch simple_divisible
+```
+
 [`log`]: https://docs.rs/log
 [`fern`]: https://docs.rs/fern
 [logging-docs]: ../user_guide.md#setting-the-log-level
@@ -117,6 +133,7 @@ expected, before you commit these updated test files.
 [test fixtures]: https://en.wikipedia.org/wiki/Test_fixture
 [`rstest`]: https://docs.rs/rstest
 [`fixture.rs`]: https://github.com/EnergySystemsModellingLab/MUSE2/blob/main/src/fixture.rs
+[patched example]: https://energysystemsmodellinglab.github.io/MUSE2/api/muse2/patch/index.html
 
 ## Example models
 
@@ -158,7 +175,7 @@ For more information, consult [the documentation for the `units` module][units-m
 Input and output files for MUSE2 are either in [CSV] or [TOML] format. Users provide model
 definitions via a number of input files and the simulation results are written to files in an output
 folder. The code responsible for reading and validating input files and writing output files is in
-the [`input`][input-module] and [`output`][output-module], respectively.
+the [`input`][input-module] and [`output`][output-module] modules, respectively.
 
 The file formats for MUSE2 input and output files are described [in the
 documentation][file-format-docs]. This documentation is generated from schema files ([JSON schemas]

justfile

@@ -19,8 +19,8 @@ coverage *ARGS:
     @cargo llvm-cov --html {{ARGS}}
 
 # Regenerate data for regression tests
-regenerate_test_data:
-    @tests/regenerate_all_data.sh
+regenerate_test_data *ARGS:
+    @tests/regenerate_test_data.sh {{ARGS}}
 
 # Run the pre-commit tool
 pre-commit *ARGS:

src/cli/example.rs

@@ -1,21 +1,24 @@
 //! Code related to the example models and the CLI commands for interacting with them.
 use super::{RunOpts, handle_run_command};
+use crate::example::patches::{get_patch_names, get_patches};
+use crate::example::{Example, get_example_names};
+use crate::patch::ModelPatch;
 use crate::settings::Settings;
-use anyhow::{Context, Result, ensure};
+use anyhow::{Context, Result};
 use clap::Subcommand;
-use include_dir::{Dir, DirEntry, include_dir};
 use std::fs;
 use std::path::{Path, PathBuf};
 use tempfile::TempDir;
 
-/// The directory containing the example models.
-const EXAMPLES_DIR: Dir = include_dir!("examples");
-
 /// The available subcommands for managing example models.
 #[derive(Subcommand)]
 pub enum ExampleSubcommands {
     /// List available examples.
-    List,
+    List {
+        /// Whether to list patched models.
+        #[arg(long, hide = true)]
+        patch: bool,
+    },
     /// Provide information about the specified example.
     Info {
         /// The name of the example.
@@ -27,11 +30,17 @@ pub enum ExampleSubcommands {
         name: String,
         /// The destination folder for the example.
         new_path: Option<PathBuf>,
+        /// Whether the model to extract is a patched model.
+        #[arg(long, hide = true)]
+        patch: bool,
     },
     /// Run an example.
     Run {
         /// The name of the example to run.
         name: String,
+        /// Whether the model to run is a patched model.
+        #[arg(long, hide = true)]
+        patch: bool,
         /// Other run options
         #[command(flatten)]
         opts: RunOpts,
@@ -42,81 +51,113 @@ impl ExampleSubcommands {
     /// Execute the supplied example subcommand
     pub fn execute(self) -> Result<()> {
         match self {
-            Self::List => handle_example_list_command(),
+            Self::List { patch } => handle_example_list_command(patch),
             Self::Info { name } => handle_example_info_command(&name)?,
             Self::Extract {
                 name,
-                new_path: dest,
-            } => handle_example_extract_command(&name, dest.as_deref())?,
-            Self::Run { name, opts } => handle_example_run_command(&name, &opts, None)?,
+                patch,
+                new_path,
+            } => handle_example_extract_command(&name, new_path.as_deref(), patch)?,
+            Self::Run { name, patch, opts } => {
+                handle_example_run_command(&name, patch, &opts, None)?;
+            }
         }
 
         Ok(())
     }
 }
 
 /// Handle the `example list` command.
-fn handle_example_list_command() {
-    for entry in EXAMPLES_DIR.dirs() {
-        println!("{}", entry.path().display());
+fn handle_example_list_command(patch: bool) {
+    if patch {
+        for name in get_patch_names() {
+            println!("{name}");
+        }
+    } else {
+        for name in get_example_names() {
+            println!("{name}");
+        }
     }
 }
 
 /// Handle the `example info` command.
 fn handle_example_info_command(name: &str) -> Result<()> {
-    let path: PathBuf = [name, "README.txt"].iter().collect();
-    let readme = EXAMPLES_DIR
-        .get_file(path)
-        .context("Example not found.")?
-        .contents_utf8()
-        .expect("README.txt is not UTF-8 encoded");
-
-    print!("{readme}");
+    // If we can't load it, it's a bug, hence why we panic
+    let info = Example::from_name(name)?
+        .get_readme()
+        .unwrap_or_else(|_| panic!("Could not load README.txt for '{name}' example"));
+    print!("{info}");
 
     Ok(())
 }
 
 /// Handle the `example extract` command
-fn handle_example_extract_command(name: &str, dest: Option<&Path>) -> Result<()> {
-    let dest = dest.unwrap_or(Path::new(name));
-    extract_example(name, dest)
+fn handle_example_extract_command(name: &str, dest: Option<&Path>, patch: bool) -> Result<()> {
+    extract_example(name, patch, dest.unwrap_or(Path::new(name)))
 }
 
-/// Extract the specified example to a new directory
-fn extract_example(name: &str, new_path: &Path) -> Result<()> {
-    // Find the subdirectory in EXAMPLES_DIR whose name matches `name`.
-    let sub_dir = EXAMPLES_DIR.get_dir(name).context("Example not found.")?;
+/// Extract the specified example to a new directory.
+///
+/// If `patch` is `true`, then the corresponding patched example will be extracted.
+fn extract_example(name: &str, patch: bool, dest: &Path) -> Result<()> {
+    if patch {
+        let patches = get_patches(name)?;
 
-    ensure!(
-        !new_path.exists(),
-        "Destination directory {} already exists",
-        new_path.display()
-    );
+        // NB: All patched models are based on `simple`, for now
+        let example = Example::from_name("simple").unwrap();
 
-    // Copy the contents of the subdirectory to the destination
-    fs::create_dir(new_path)?;
-    for entry in sub_dir.entries() {
-        match entry {
-            DirEntry::Dir(_) => panic!("Subdirectories in examples not supported"),
-            DirEntry::File(f) => {
-                let file_name = f.path().file_name().unwrap();
-                let file_path = new_path.join(file_name);
-                fs::write(&file_path, f.contents())?;
-            }
-        }
-    }
+        // First extract the example to a temp dir
+        let example_tmp = TempDir::new().context("Failed to create temporary directory")?;
+        let example_path = example_tmp.path().join(name);
+        example
+            .extract(&example_path)
+            .context("Could not extract example")?;
 
-    Ok(())
+        // Patch example and put contents in dest
+        fs::create_dir(dest).context("Could not create output directory")?;
+        ModelPatch::new(example_path)
+            .with_file_patches(patches.to_owned())
+            .build(dest)
+            .context("Failed to patch example")
+    } else {
+        // Otherwise it's just a regular example
+        let example = Example::from_name(name)?;
+        example.extract(dest)
+    }
 }
 
 /// Handle the `example run` command.
 pub fn handle_example_run_command(
     name: &str,
+    patch: bool,
     opts: &RunOpts,
     settings: Option<Settings>,
 ) -> Result<()> {
-    let temp_dir = TempDir::new().context("Failed to create temporary directory.")?;
+    let temp_dir = TempDir::new().context("Failed to create temporary directory")?;
     let model_path = temp_dir.path().join(name);
-    extract_example(name, &model_path)?;
+    extract_example(name, patch, &model_path)?;
     handle_run_command(&model_path, opts, settings)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    fn assert_dir_non_empty(path: &Path) {
+        assert!(
+            path.read_dir().unwrap().next().is_some(),
+            "Directory is empty"
+        );
+    }
+
+    #[rstest]
+    #[case("muse1_default", false)]
+    #[case("simple_divisible", true)]
+    fn check_extract_example(#[case] name: &str, #[case] patch: bool) {
+        let tmp = TempDir::new().unwrap();
+        let dest = tmp.path().join("out");
+        extract_example(name, patch, &dest).unwrap();
+        assert_dir_non_empty(&dest);
+    }
+}

src/example.rs

@@ -0,0 +1,82 @@
+//! Code for working with example models
+use std::fs;
+use std::path::Path;
+
+use anyhow::{Context, Result};
+use include_dir::{Dir, DirEntry, include_dir};
+
+pub mod patches;
+
+/// The directory containing the example models.
+const EXAMPLES_DIR: Dir = include_dir!("examples");
+
+/// Get the names of all examples
+pub fn get_example_names() -> impl Iterator<Item = &'static str> {
+    EXAMPLES_DIR.dirs().map(|dir| {
+        dir.path()
+            .as_os_str()
+            .to_str()
+            .expect("Invalid unicode in path")
+    })
+}
+
+/// A bundled example model
+pub struct Example(Dir<'static>);
+
+impl Example {
+    /// Get the example with the specified name
+    pub fn from_name(name: &str) -> Result<Self> {
+        let dir = EXAMPLES_DIR
+            .get_dir(name)
+            .with_context(|| format!("Example '{name}' not found"))?;
+
+        Ok(Self(dir.clone()))
+    }
+
+    /// Get the contents of the readme file for this example
+    pub fn get_readme(&self) -> Result<&'static str> {
+        self.0
+            .get_file(self.0.path().join("README.txt"))
+            .context("Missing file")?
+            .contents_utf8()
+            .context("File not UTF-8 encoded")
+    }
+
+    /// Extract this example to a specified destination.
+    ///
+    /// Returns an error if the destination directory already exists or copying the files fails.
+    pub fn extract(&self, new_path: &Path) -> Result<()> {
+        // Copy the contents of the subdirectory to the destination
+        fs::create_dir(new_path)?;
+        for entry in self.0.entries() {
+            match entry {
+                DirEntry::Dir(_) => panic!("Subdirectories in examples not supported"),
+                DirEntry::File(f) => {
+                    let file_name = f.path().file_name().unwrap();
+                    let file_path = new_path.join(file_name);
+                    fs::write(&file_path, f.contents())?;
+                }
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn all_examples_have_readme() {
+        for example in get_example_names() {
+            let readme = Example::from_name(example)
+                .unwrap()
+                .get_readme()
+                .with_context(|| format!("Could not load readme for {example}"))
+                .unwrap();
+
+            assert!(!readme.trim().is_empty());
+        }
+    }
+}

src/example/patches.rs

@@ -0,0 +1,39 @@
+//! File patches to be used in integration tests.
+//!
+//! This is used to test small variations on existing example models.
+use crate::patch::FilePatch;
+use anyhow::{Context, Result};
+use std::{collections::BTreeMap, sync::LazyLock};
+
+/// A map of file patches, keyed by name
+type PatchMap = BTreeMap<&'static str, Vec<FilePatch>>;
+
+/// The file patches, keyed by name
+static PATCHES: LazyLock<PatchMap> = LazyLock::new(get_all_patches);
+
+/// Get all patches
+fn get_all_patches() -> PatchMap {
+    [(
+        // The simple example with gas boiler process made divisible
+        "simple_divisible",
+        vec![
+            FilePatch::new("processes.csv")
+                .with_deletion("RGASBR,Gas boiler,all,RSHEAT,2020,2040,1.0,")
+                .with_addition("RGASBR,Gas boiler,all,RSHEAT,2020,2040,1.0,1000"),
+        ],
+    )]
+    .into_iter()
+    .collect()
+}
+
+/// Get the names for all the patches
+pub fn get_patch_names() -> impl Iterator<Item = &'static str> {
+    PATCHES.keys().copied()
+}
+
+/// Get patches for the named patched example
+pub fn get_patches(name: &str) -> Result<&[FilePatch]> {
+    Ok(PATCHES
+        .get(name)
+        .with_context(|| format!("Patched example '{name}' not found"))?)
+}

src/lib.rs

@@ -11,6 +11,7 @@ pub mod agent;
 pub mod asset;
 pub mod cli;
 pub mod commodity;
+pub mod example;
 pub mod finance;
 pub mod graph;
 pub mod id;