From daeec327ba992b76b61cf9774615adbbada1c5eb Mon Sep 17 00:00:00 2001
From: HaYeong Lee <110161434+glorialeezero@users.noreply.github.com>
Date: Mon, 25 Nov 2024 13:55:00 +0900
Subject: [PATCH] Add testcase for #10

---
 tests/test_transition.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tests/test_transition.py b/tests/test_transition.py
index 21010ec..4d34822 100644
--- 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