From 4635273c1e1f4af38471b58666ae7e5518fd1b92 Mon Sep 17 00:00:00 2001
From: Agustin Borgna <agustin.borgna@quantinuum.com>
Date: Fri, 28 Jun 2024 11:26:19 +0100
Subject: [PATCH] Add a flat test for SiblingGraph in addition to the
 DescendantGraph one

---
 hugr-core/src/hugr/views/descendants.rs | 19 +++++++++++++++++++
 hugr-core/src/hugr/views/sibling.rs     |  7 +++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/hugr-core/src/hugr/views/descendants.rs b/hugr-core/src/hugr/views/descendants.rs
index 35a6e46c6..d2b66b250 100644
--- 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
index bdc72ca1c..7131c2451 100644
--- 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());