Add testcase for #10

kupl · Nov 25, 2024 · daeec32 · daeec32
1 parent 0872367
commit daeec32
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/tests/test_transition.py b/tests/test_transition.py
@@ -2,6 +2,7 @@
     ALL_AEXPS,
     ALL_GATES,
     CX,
+    Add,
     Aexp,
     Cmd,
     ForCmd,
@@ -78,3 +79,17 @@ def test_hole_aexp3():
         if type(pgm.body.gate.qreg2) not in [Integer, Var]:
             aexp_types.remove(type(pgm.body.gate.qreg2))
     assert len(aexp_types) == 3
+
+
+def test_next_aexp():
+    pgm = Pgm("n", GateCmd(CX(Integer(0), Add())))
+    aexp_types: list[type[Aexp]] = ALL_AEXPS.copy()
+    pgms = next(pgm)
+    for pgm in pgms:
+        assert isinstance(pgm.body, GateCmd)
+        assert isinstance(pgm.body.gate, CX)
+        assert isinstance(pgm.body.gate.qreg2, Add)
+        assert type(pgm.body.gate.qreg2.a) in aexp_types
+        if type(pgm.body.gate.qreg2.a) not in [Integer, Var]:
+            aexp_types.remove(type(pgm.body.gate.qreg2.a))
+    assert len(aexp_types) == 3