1 files changed, 97 insertions, 0 deletions
diff --git a/tvix/glue/src/known_paths.rs b/tvix/glue/src/known_paths.rs
index 1e3fe85a0f7c..03b04d4fc85f 100644
--- a/tvix/glue/src/known_paths.rs
+++ b/tvix/glue/src/known_paths.rs
@@ -108,3 +108,100 @@ impl KnownPaths {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use nix_compat::{derivation::Derivation, nixhash::NixHash, store_path::StorePath};
+
+    use super::KnownPaths;
+    use hex_literal::hex;
+    use lazy_static::lazy_static;
+
+    lazy_static! {
+        static ref BAR_DRV: Derivation = Derivation::from_aterm_bytes(include_bytes!(
+            "tests/ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv"
+        ))
+        .expect("must parse");
+        static ref FOO_DRV: Derivation = Derivation::from_aterm_bytes(include_bytes!(
+            "tests/ch49594n9avinrf8ip0aslidkc4lxkqv-foo.drv"
+        ))
+        .expect("must parse");
+        static ref BAR_DRV_PATH: StorePath =
+            StorePath::from_bytes(b"ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv").expect("must parse");
+        static ref FOO_DRV_PATH: StorePath =
+            StorePath::from_bytes(b"ch49594n9avinrf8ip0aslidkc4lxkqv-foo.drv").expect("must parse");
+        static ref BAR_OUT_PATH: StorePath =
+            StorePath::from_bytes(b"mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar").expect("must parse");
+        static ref FOO_OUT_PATH: StorePath =
+            StorePath::from_bytes(b"fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo").expect("must parse");
+    }
+
+    /// ensure we don't allow acdding a Derivation that depends on another,
+    /// not-yet-added Derivation.
+    #[test]
+    #[should_panic]
+    fn reject_if_missing_input_drv() {
+        let mut known_paths = KnownPaths::default();
+
+        // FOO_DRV depends on BAR_DRV, which wasn't added.
+        known_paths.add(FOO_DRV_PATH.clone(), FOO_DRV.clone());
+    }
+
+    #[test]
+    fn happy_path() {
+        let mut known_paths = KnownPaths::default();
+
+        // get_drv_by_drvpath should return None for non-existing Derivations,
+        // same as get_hash_derivation_modulo and get_drv_path_for_output_path
+        assert_eq!(None, known_paths.get_drv_by_drvpath(&BAR_DRV_PATH));
+        assert_eq!(None, known_paths.get_hash_derivation_modulo(&BAR_DRV_PATH));
+        assert_eq!(
+            None,
+            known_paths.get_drv_path_for_output_path(&BAR_OUT_PATH)
+        );
+
+        // Add BAR_DRV
+        known_paths.add(BAR_DRV_PATH.clone(), BAR_DRV.clone());
+
+        // We should get it back
+        assert_eq!(
+            Some(&BAR_DRV.clone()),
+            known_paths.get_drv_by_drvpath(&BAR_DRV_PATH)
+        );
+
+        // Test get_drv_path_for_output_path
+        assert_eq!(
+            Some(&BAR_DRV_PATH.clone()),
+            known_paths.get_drv_path_for_output_path(&BAR_OUT_PATH)
+        );
+
+        // It should be possible to get the hash derivation modulo.
+        assert_eq!(
+            Some(&NixHash::Sha256(hex!(
+                "c79aebd0ce3269393d4a1fde2cbd1d975d879b40f0bf40a48f550edc107fd5df"
+            ))),
+            known_paths.get_hash_derivation_modulo(&BAR_DRV_PATH.clone())
+        );
+
+        // Now insert FOO_DRV too. It shouldn't panic, as BAR_DRV is already
+        // added.
+        known_paths.add(FOO_DRV_PATH.clone(), FOO_DRV.clone());
+
+        assert_eq!(
+            Some(&FOO_DRV.clone()),
+            known_paths.get_drv_by_drvpath(&FOO_DRV_PATH)
+        );
+        assert_eq!(
+            Some(&NixHash::Sha256(hex!(
+                "af030d36d63d3d7f56a71adaba26b36f5fa1f9847da5eed953ed62e18192762f"
+            ))),
+            known_paths.get_hash_derivation_modulo(&FOO_DRV_PATH.clone())
+        );
+
+        // Test get_drv_path_for_output_path
+        assert_eq!(
+            Some(&FOO_DRV_PATH.clone()),
+            known_paths.get_drv_path_for_output_path(&FOO_OUT_PATH)
+        );
+    }
+}