Improve test

CQCL · Dec 11, 2024 · 36ad0a0 · 36ad0a0
1 parent e1ac29a
commit 36ad0a0
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/hugr-core/src/extension/prelude/generic.rs b/hugr-core/src/extension/prelude/generic.rs
@@ -167,7 +167,7 @@ mod tests {
         ops::{constant, OpType},
         type_row,
         types::TypeArg,
-        OutgoingPort,
+        HugrView, OutgoingPort,
     };
 
     use super::LoadNat;
@@ -179,9 +179,19 @@ mod tests {
         let arg = TypeArg::BoundedNat { n: 4 };
         let op = LoadNat::new(arg);
 
-        let out = b.add_dataflow_op(op, []).unwrap();
+        let out = b.add_dataflow_op(op.clone(), []).unwrap();
 
-        b.finish_hugr_with_outputs(out.outputs()).unwrap();
+        let result = b.finish_hugr_with_outputs(out.outputs()).unwrap();
+
+        let exp_optype: OpType = op.into();
+
+        for child in result.children(result.root()) {
+            let node_optype = result.get_optype(child);
+            // The only node in the HUGR besides Input and Output should be LoadNat.
+            if !node_optype.is_input() && !node_optype.is_output() {
+                assert_eq!(node_optype, &exp_optype)
+            }
+        }
     }
 
     #[test]