Add a flat test for SiblingGraph in addition to the DescendantGraph one

CQCL · Jun 28, 2024 · 4635273 · 4635273
1 parent f5df394
commit 4635273
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/hugr-core/src/hugr/views/descendants.rs b/hugr-core/src/hugr/views/descendants.rs
@@ -201,6 +201,9 @@ where
 
 #[cfg(test)]
 pub(super) mod test {
+    use rstest::rstest;
+
+    use crate::extension::PRELUDE_REGISTRY;
     use crate::{
         builder::{Container, Dataflow, DataflowSubContainer, HugrBuilder, ModuleBuilder},
         type_row,
@@ -269,4 +272,20 @@ pub(super) mod test {
 
         Ok(())
     }
+
+    #[rstest]
+    fn extract_hugr() -> Result<(), Box<dyn std::error::Error>> {
+        let (hugr, def, _inner) = make_module_hgr()?;
+
+        let region: DescendantsGraph = DescendantsGraph::try_new(&hugr, def)?;
+        let extracted = region.extract_hugr();
+        extracted.validate(&PRELUDE_REGISTRY)?;
+
+        let region: DescendantsGraph = DescendantsGraph::try_new(&hugr, def)?;
+
+        assert_eq!(region.node_count(), extracted.node_count());
+        assert_eq!(region.root_type(), extracted.root_type());
+
+        Ok(())
+    }
 }
diff --git a/hugr-core/src/hugr/views/sibling.rs b/hugr-core/src/hugr/views/sibling.rs
@@ -381,7 +381,6 @@ mod test {
     use crate::builder::test::simple_dfg_hugr;
     use crate::builder::{Container, Dataflow, DataflowSubContainer, HugrBuilder, ModuleBuilder};
     use crate::extension::PRELUDE_REGISTRY;
-    use crate::hugr::views::DescendantsGraph;
     use crate::ops::handle::{CfgID, DataflowParentID, DfgID, FuncID};
     use crate::ops::{dataflow::IOTrait, Input, OpTag, Output};
     use crate::ops::{OpTrait, OpType};
@@ -492,13 +491,13 @@ mod test {
 
     #[rstest]
     fn extract_hugr() -> Result<(), Box<dyn std::error::Error>> {
-        let (hugr, def, _inner) = make_module_hgr()?;
+        let (hugr, _def, inner) = make_module_hgr()?;
 
-        let region: DescendantsGraph = DescendantsGraph::try_new(&hugr, def)?;
+        let region: SiblingGraph = SiblingGraph::try_new(&hugr, inner)?;
         let extracted = region.extract_hugr();
         extracted.validate(&PRELUDE_REGISTRY)?;
 
-        let region: DescendantsGraph = DescendantsGraph::try_new(&hugr, def)?;
+        let region: SiblingGraph = SiblingGraph::try_new(&hugr, inner)?;
 
         assert_eq!(region.node_count(), extracted.node_count());
         assert_eq!(region.root_type(), extracted.root_type());