Fix test because of par_domain removal (#88)

* Fix test par_domain

* Comment from review.

tests/test_numpy_wavefunction.py

@@ -70,65 +70,56 @@ def test_var_hermitian(self, tol):
         self.assertAlmostEqual(var, expected, delta=tol)
 
     @pytest.mark.parametrize(
-        "gate", plf.NumpyWavefunctionDevice._operation_map
-    )  # pylint: disable=protected-access
-    def test_apply(self, gate, apply_unitary, tol):
+        "op",
+        [
+            qml.QubitUnitary(np.array(U), wires=0),
+            qml.BasisState(np.array([1, 1, 1]), wires=list(range(3))),
+            qml.PauliX(wires=0),
+            qml.PauliY(wires=0),
+            qml.PauliZ(wires=0),
+            qml.S(wires=0),
+            qml.T(wires=0),
+            qml.RX(0.432, wires=0),
+            qml.RY(0.432, wires=0),
+            qml.RZ(0.432, wires=0),
+            qml.Hadamard(wires=0),
+            qml.Rot(0.432, 2, 0.324, wires=0),
+            qml.Toffoli(wires=[0, 1, 2]),
+            qml.SWAP(wires=[0, 1]),
+            qml.CSWAP(wires=[0, 1, 2]),
+            qml.CZ(wires=[0, 1]),
+            qml.CNOT(wires=[0, 1]),
+            qml.PhaseShift(0.432, wires=0),
+            qml.CSWAP(wires=[0, 1, 2]),
+            plf.CPHASE(0.432, 2, wires=[0, 1]),
+            plf.ISWAP(wires=[0, 1]),
+            plf.PSWAP(0.432, wires=[0, 1]),
+        ],
+    )
+    def test_apply(self, op, apply_unitary, tol):
         """Test the application of gates to a state"""
         dev = plf.NumpyWavefunctionDevice(wires=3)
 
-        try:
-            # get the equivalent pennylane operation class
-            op = getattr(qml.ops, gate)
-        except AttributeError:
-            # get the equivalent pennylane-forest operation class
-            op = getattr(plf, gate)
-
-        # the list of wires to apply the operation to
-        w = list(range(op.num_wires))
-
         obs = qml.expval(qml.PauliZ(0))
-        if op.par_domain == "A":
-            # the parameter is an array
-            if gate == "QubitUnitary":
-                p = np.array(U)
-                w = [0]
-                state = apply_unitary(U, 3)
-            elif gate == "BasisState":
-                p = np.array([1, 1, 1])
-                state = np.array([0, 0, 0, 0, 0, 0, 0, 1])
-                w = list(range(dev.num_wires))
-
-            with qml.tape.QuantumTape() as tape:
-                op(p, wires=w)
-                obs
+
+        if op.name == "QubitUnitary":
+            state = apply_unitary(U, 3)
+        elif op.name == "BasisState":
+            state = np.array([0, 0, 0, 0, 0, 0, 0, 1])
+        elif op.name == "CPHASE":
+            state = apply_unitary(test_operation_map["CPHASE"](0.432, 2), 3)
+        elif op.name == "ISWAP":
+            state = apply_unitary(test_operation_map["ISWAP"], 3)
+        elif op.name == "PSWAP":
+            state = apply_unitary(test_operation_map["PSWAP"](0.432), 3)
         else:
-            p = [0.432_423, 2, 0.324][: op.num_params]
-            fn = test_operation_map[gate]
-            if callable(fn):
-                # if the default.qubit is an operation accepting parameters,
-                # initialise it using the parameters generated above.
-                O = fn(*p)
-            else:
-                # otherwise, the operation is simply an array.
-                O = fn
-
-            # calculate the expected output
-            state = apply_unitary(O, 3)
-            # Creating the tape using a parametrized operation
-            if p:
-                with qml.tape.QuantumTape() as tape:
-                    op(*p, wires=w)
-                    obs
-
-            # Creating the tape using an operation that take no parameters
-            else:
-                with qml.tape.QuantumTape() as tape:
-                    op(wires=w)
-                    obs
+            state = apply_unitary(op.matrix, 3)
 
-        dev.apply(tape.operations, rotations=tape.diagonalizing_gates)
+        with qml.tape.QuantumTape() as tape:
+            qml.apply(op)
+            obs
 
-        res = dev.expval(obs)
+        dev.apply(tape.operations, rotations=tape.diagonalizing_gates)
 
         # verify the device is now in the expected state
         self.assertAllAlmostEqual(dev._state, state, delta=tol)
@@ -402,5 +393,4 @@ def circuit(x, y, z):
             runs.append(circuit(a, b, c))
 
         expected_var = np.sqrt(1 / shots)
-        print(np.mean(runs), np.cos(a) * np.sin(b))
         self.assertAlmostEqual(np.mean(runs), np.cos(a) * np.sin(b), delta=expected_var)