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

---
 mitiq/lre/lre.py            |  9 +------
 mitiq/lre/tests/test_lre.py | 51 +++++++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 13 deletions(-)

diff --git a/mitiq/lre/lre.py b/mitiq/lre/lre.py
index dd0354897..32ce2cf86 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
     )
@@ -89,13 +89,6 @@ def lre_decorator(
 ) -> Callable[
     [Callable[[Circuit], QuantumResult]], Callable[[Circuit], float]
 ]:
-    # Raise an error if the decorator is used without parenthesis
-    if callable(observable):
-        raise TypeError(
-            "Decorator must be used with parentheses (i.e., @lre_decorator()) "
-            "with explicit arguments for shots, degree and fold_multiplier."
-        )
-
     def decorator(
         executor: Callable[[Circuit], QuantumResult],
     ) -> Callable[[Circuit], float]:
diff --git a/mitiq/lre/tests/test_lre.py b/mitiq/lre/tests/test_lre.py
index e4ef5c4e9..93addfb91 100644
--- a/mitiq/lre/tests/test_lre.py
+++ b/mitiq/lre/tests/test_lre.py
@@ -1,9 +1,12 @@
 """Unit tests for the LRE extrapolation methods."""
 
+import re
+
+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 +15,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 +22,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, match=re.escape("lre_decorator() missing")):
+
+        @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)