From a399d36b41d85331a61576600ae1c89a0ea8426a Mon Sep 17 00:00:00 2001
From: Hosseinberg <34165958+Hosseinberg@users.noreply.github.com>
Date: Tue, 25 Apr 2023 02:12:02 -0400
Subject: [PATCH] adds test for _controlled_qubit_unitary

---
 tests/test_apply.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tests/test_apply.py b/tests/test_apply.py
index a8991f0..000d440 100755
--- a/tests/test_apply.py
+++ b/tests/test_apply.py
@@ -249,6 +249,23 @@ def test_qubit_unitary(self, init_state, mat, tol):
         expected = mat @ state
         assert np.allclose(res, expected, tol)
 
+    @pytest.mark.parametrize("control_wires", [[0], [1], [0, 1]])
+    @pytest.mark.parametrize("mat", [U, U2])
+    def test_controlled_qubit_unitary(self, init_state, mat, control_wires, tol):
+        """Test ControlledQubitUnitary application"""
+
+        N = int(np.log2(len(mat)))
+        dev = QulacsDevice(N + len(control_wires))
+        state = init_state(N + len(control_wires))
+
+        op = qml.ControlledQubitUnitary(mat, wires=list(range(N, N + len(control_wires) + 1)), control_wires=control_wires)
+        dev.apply([qml.QubitStateVector(state, wires=list(range(N + len(control_wires)))), op])
+        dev._obs_queue = []
+
+        res = dev.state
+        expected = np.kron(np.eye(2 ** N), np.eye(2 ** len(control_wires))) @ np.kron(np.eye(2 ** N), mat) @ state
+        assert np.allclose(res, expected, tol)
+
     def test_invalid_qubit_state_unitary(self):
         """Test that an exception is raised if the
         unitary matrix is the wrong size"""