From c0ff6ea5e72a71221ba23f40af8b7a2b0db25ed4 Mon Sep 17 00:00:00 2001
From: David Seddon <david.seddon@octoenergy.com>
Date: Fri, 4 Oct 2024 17:20:29 +0100
Subject: [PATCH] Fix flakey test by turning get_descendants into Result

---
 rust/src/graph.rs               | 46 +++++++++++++++++++++++++--------
 rust/src/lib.rs                 |  2 +-
 src/grimp/adaptors/rustgraph.py |  1 +
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/rust/src/graph.rs b/rust/src/graph.rs
index e10a93c8..45a69df6 100644
--- a/rust/src/graph.rs
+++ b/rust/src/graph.rs
@@ -47,6 +47,17 @@ impl fmt::Display for Module {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct ModuleNotPresent{
+    pub module: Module
+}
+
+
+impl fmt::Display for ModuleNotPresent {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "\"{}\" not present in the graph", self.module.name)
+    }
+}
 
 impl Module {
     pub fn new(name: String) -> Module {
@@ -185,13 +196,16 @@ impl Graph {
             .collect()
     }
 
-    pub fn find_descendants(&self, module: &Module) -> HashSet<&Module> {
-        let module_index = self.hierarchy_module_indices.get_by_left(module).unwrap();
-        Bfs::new(&self.hierarchy, *module_index)
+    pub fn find_descendants(&self, module: &Module) -> Result<HashSet<&Module>, ModuleNotPresent> {
+        let module_index = match self.hierarchy_module_indices.get_by_left(module) {
+            Some(index) => index,
+            None => return Err(ModuleNotPresent{module: module.clone()}),
+        };
+        Ok(Bfs::new(&self.hierarchy, *module_index)
             .iter(&self.hierarchy)
             .filter(|index| index != module_index) // Don't include the supplied module.
-            .map(|index| self.hierarchy_module_indices.get_by_right(&index).unwrap())
-            .collect()
+            .map(|index| self.hierarchy_module_indices.get_by_right(&index).unwrap())  // This panics sometimes.
+            .collect())
     }
 
     pub fn add_import(&mut self, importer: &Module, imported: &Module) {
@@ -420,13 +434,13 @@ impl Graph {
 
         let mut importer_modules: HashSet<&Module> = HashSet::from([importer]);
         // TODO don't do this if module is squashed?
-        for descendant in self.find_descendants(&importer) {
+        for descendant in self.find_descendants(&importer).unwrap() {
             importer_modules.insert(descendant);
         }
 
         let mut imported_modules: HashSet<&Module> = HashSet::from([imported]);
         // TODO don't do this if module is squashed?
-        for descendant in self.find_descendants(&imported) {
+        for descendant in self.find_descendants(&imported).unwrap() {
             imported_modules.insert(descendant);
         }
 
@@ -465,7 +479,7 @@ impl Graph {
     #[allow(unused_variables)]
     pub fn squash_module(&mut self, module: &Module) {
         // Get descendants and their imports.
-        let descendants: Vec<Module> = self.find_descendants(module).into_iter().cloned().collect();
+        let descendants: Vec<Module> = self.find_descendants(module).unwrap().into_iter().cloned().collect();
         let modules_imported_by_descendants: Vec<Module> = descendants
             .iter()
             .flat_map(|descendant| {
@@ -776,7 +790,17 @@ imports:
         graph.add_module(mypackage_foo_alpha_green.clone());
         graph.add_module(mypackage_foo_beta.clone());
 
-        assert_eq!(graph.find_descendants(&mypackage_bar), HashSet::new());
+        assert_eq!(graph.find_descendants(&mypackage_bar), Ok(HashSet::new()));
+    }
+
+    #[test]
+    fn find_descendants_module_not_in_graph() {
+        let mut graph = Graph::default();
+        let blue = Module::new("blue".to_string());
+        let green = Module::new("green".to_string());
+        graph.add_module(blue.clone());
+
+        assert_eq!(graph.find_descendants(&green), Err(ModuleNotPresent{module: green.clone()}));
     }
 
     #[test]
@@ -800,12 +824,12 @@ imports:
 
         assert_eq!(
             graph.find_descendants(&mypackage_foo),
-            HashSet::from([
+            Ok(HashSet::from([
                 &mypackage_foo_alpha,
                 &mypackage_foo_alpha_blue,
                 &mypackage_foo_alpha_green,
                 &mypackage_foo_beta
-            ])
+            ]))
         );
     }
 
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 111303a9..29e4beed 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -114,7 +114,7 @@ impl GraphWrapper {
 
     pub fn find_descendants(&self, module: &str) -> HashSet<String> {
         self._graph
-            .find_descendants(&Module::new(module.to_string()))
+            .find_descendants(&Module::new(module.to_string())).unwrap()
             .iter()
             .map(|descendant| descendant.name.clone())
             .collect()
diff --git a/src/grimp/adaptors/rustgraph.py b/src/grimp/adaptors/rustgraph.py
index 8b09e85d..04a97218 100644
--- a/src/grimp/adaptors/rustgraph.py
+++ b/src/grimp/adaptors/rustgraph.py
@@ -133,4 +133,5 @@ def find_illegal_dependencies_for_layers(
         layers: Sequence[Layer | str | set[str]],
         containers: set[str] | None = None,
     ) -> set[PackageDependency]:
+        # TODO
         return super().find_illegal_dependencies_for_layers(layers, containers)