From e328b9d71b85d5d35abff33beb2bacb7fe64b263 Mon Sep 17 00:00:00 2001
From: Purva Thakre <purva.athakre@gmail.com>
Date: Wed, 11 Sep 2024 14:01:41 -0500
Subject: [PATCH] add unit tests for the decorator - check for test coverage

---
 mitiq/lre/lre.py            |  2 +-
 mitiq/lre/tests/test_lre.py | 49 +++++++++++++++++++++++++++++++++----
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index dd0354897..8b943b588 100644
--- a/mitiq/lre/lre.py
+++ b/mitiq/lre/lre.py
@@ -28,7 +28,7 @@ def execute_with_lre(
     folding_method: Callable[[Circuit, float], Circuit] = fold_gates_at_random,
     num_chunks: Optional[int] = None,
     observable: Optional[Observable] = None,
-)-> float:
+) -> float:
     noise_scaled_circuits = multivariate_layer_scaling(
         input_circuit, degree, fold_multiplier, num_chunks, folding_method
     )
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index e4ef5c4e9..155ebd643 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -1,9 +1,10 @@
 """Unit tests for the LRE extrapolation methods."""
 
+import pytest
 from cirq import DensityMatrixSimulator, depolarize
 
 from mitiq import benchmarks
-from mitiq.lre import execute_with_lre
+from mitiq.lre import execute_with_lre, lre_decorator, mitigate_executor
 
 test_cirq = benchmarks.generate_rb_circuits(
     n_qubits=1,
@@ -12,9 +13,6 @@
 
 
 def execute(circuit, noise_level=0.025, shots=1000):
-    """Returns Tr[ρ |0⟩⟨0|] where ρ is the state prepared by the circuit
-    executed with depolarizing noise.
-    """
     # Replace with code based on your frontend and backend.
     mitiq_circuit = circuit
     noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
@@ -22,9 +20,50 @@ def execute(circuit, noise_level=0.025, shots=1000):
     return rho[0, 0].real
 
 
+noisy_val = execute(test_cirq)
+
+
 def test_lre_exp_value():
-    noisy_val = execute(test_cirq)
     ideal_val = execute(test_cirq, noise_level=0, shots=1000)
     assert abs(ideal_val - noisy_val) > 0
     lre_exp_val = execute_with_lre(test_cirq, execute, 1000, 2, 2)
     assert lre_exp_val > noisy_val
+
+    # verify the mitigated decorator work as expected
+    mitigated_executor = mitigate_executor(execute, 1000, 2, 2)
+    exp_val_from_mitigate_executor = mitigated_executor(test_cirq)
+    assert exp_val_from_mitigate_executor > noisy_val
+
+
+def test_lre_decorator():
+    @lre_decorator(100, 2, 2)
+    def execute(circuit, noise_level=0.025, shots=100):
+        # Replace with code based on your frontend and backend.
+        mitiq_circuit = circuit
+        noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+        rho = (
+            DensityMatrixSimulator()
+            .simulate(noisy_circuit)
+            .final_density_matrix
+        )
+        return rho[0, 0].real
+
+    assert noisy_val < execute(test_cirq)
+
+
+def test_lre_decorator_raised_error():
+    with pytest.raises(TypeError):
+
+        @lre_decorator()
+        def execute(circuit, noise_level=0.025, shots=100):
+            # Replace with code based on your frontend and backend.
+            mitiq_circuit = circuit
+            noisy_circuit = mitiq_circuit.with_noise(depolarize(p=noise_level))
+            rho = (
+                DensityMatrixSimulator()
+                .simulate(noisy_circuit)
+                .final_density_matrix
+            )
+            return rho[0, 0].real
+
+        assert noisy_val < execute(test_cirq)