From a7fe64622c24fdd6157350e5bc61c323310a215e Mon Sep 17 00:00:00 2001
From: king-p3nguin <kazukitsuoka@g.ecc.u-tokyo.ac.jp>
Date: Sat, 8 Jun 2024 04:30:39 +0900
Subject: [PATCH] test: add tests

---
 test/integ_tests/test_sample.py | 80 +++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/test/integ_tests/test_sample.py b/test/integ_tests/test_sample.py
index db1f8c4b..e73c48e4 100755
--- a/test/integ_tests/test_sample.py
+++ b/test/integ_tests/test_sample.py
@@ -24,6 +24,86 @@
 class TestSample:
     """Tests for the sample return type"""
 
+    def test_sample_default(self, device, shots, tol):
+        """Tests if the samples returned by sample have
+        the correct values when no observable is specified
+        """
+        dev = device(2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(np.pi / 4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.sample()
+
+        shot_vector = circuit()
+
+        # The sample should only contain 1 and 0
+        assert shot_vector.dtype == np.dtype("int")
+        assert np.all((shot_vector == 0) | (shot_vector == 1))
+        assert shot_vector.shape == (shots, 2)
+
+    def test_sample_wires(self, device, shots, tol):
+        """Tests if the samples returned by sample have
+        the correct values when only wires are specified
+        """
+        dev = device(3)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(np.pi / 4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            return qml.sample(wires=[0, 2])
+
+        shot_vector = circuit()
+
+        # The sample should only contain 1 and 0
+        assert shot_vector.dtype == np.dtype("int")
+        assert np.all((shot_vector == 0) | (shot_vector == 1))
+        assert shot_vector.shape == (shots, 2)
+
+    def test_sample_batch_default(self, device, shots, tol):
+        """Tests if the samples returned by sample have
+        the correct values when no observable is specified and
+        the batch dimension is returned
+        """
+        dev = device(3)
+
+        @qml.qnode(dev)
+        def circuit(a):
+            qml.RX(a, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            return qml.sample()
+
+        shot_vector = circuit([np.pi / 4, np.pi / 3])
+
+        # The sample should only contain 1 and 0
+        assert shot_vector.dtype == np.dtype("int")
+        assert np.all((shot_vector == 0) | (shot_vector == 1))
+        assert shot_vector.shape == (2, shots, 3)
+
+    def test_sample_batch_wires(self, device, shots, tol):
+        """Tests if the samples returned by sample have
+        the correct values when only wires are specified
+        """
+        dev = device(4)
+
+        @qml.qnode(dev)
+        def circuit(a):
+            qml.RX(a, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            return qml.sample(wires=[0, 2, 3])
+
+        shot_vector = circuit([np.pi / 4, np.pi / 3])
+
+        # The sample should only contain 1 and 0
+        assert shot_vector.dtype == np.dtype("int")
+        assert np.all((shot_vector == 0) | (shot_vector == 1))
+        assert shot_vector.shape == (2, shots, 3)
+
     def test_sample_values(self, device, shots, tol):
         """Tests if the samples returned by sample have
         the correct values