fixing test

Author: weinbe58

test/squin/test_measure_sugar.py

@@ -1,16 +1,24 @@
+from kirin import ir
+from kirin.dialects import func
+
 from bloqade import squin
 
 
+def get_return_value_stmt(kernel: ir.Method):
+    assert isinstance(
+        last_stmt := kernel.callable_region.blocks[-1].last_stmt, func.Return
+    )
+    return last_stmt.value.owner
+
+
 def test_measure_register():
     @squin.kernel
     def test_measure_sugar():
         q = squin.qubit.new(2)
 
         return squin.qubit.measure(q)
 
-    assert isinstance(
-        test_measure_sugar.callable_region.blocks[-1].last_stmt, squin.qubit.MeasureReg
-    )
+    assert isinstance(get_return_value_stmt(test_measure_sugar), squin.qubit.MeasureReg)
 
 
 def test_measure_qubit():
@@ -21,6 +29,6 @@ def test_measure_sugar():
         return squin.qubit.measure(q[0])
 
     assert isinstance(
-        test_measure_sugar.callable_region.blocks[-1].last_stmt,
+        get_return_value_stmt(test_measure_sugar),
         squin.qubit.MeasureQubit,
     )