update tests

quantumlib · Dec 20, 2024 · 1f98177 · 1f98177
1 parent 044db6e
commit 1f98177
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/cirq-core/cirq/transformers/gauge_compiling/cz_gauge_test.py b/cirq-core/cirq/transformers/gauge_compiling/cz_gauge_test.py
@@ -21,4 +21,4 @@
 class TestCZGauge(GaugeTester):
     two_qubit_gate = cirq.CZ
     gauge_transformer = CZGaugeTransformer
-    enable_test_sweep = True
+    sweep_must_pass = True
diff --git a/cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test_utils.py b/cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test_utils.py
@@ -27,7 +27,7 @@ class GaugeTester:
     two_qubit_gate: cirq.Gate
     gauge_transformer: GaugeTransformer
     must_fail: bool = False
-    enable_test_sweep: bool = False
+    sweep_must_pass: bool = False
 
     @pytest.mark.parametrize(
         ['generation_seed', 'transformation_seed'],
@@ -88,9 +88,7 @@ def test_sweep(self):
             cirq.Moment([cirq.measure(q) for q in [a, b, c]]),
         )
 
-        supported_gates: set[cirq.Gate] = {cirq.CZ}
-
-        if self.two_qubit_gate not in supported_gates:
+        if not self.sweep_must_pass:
             with pytest.raises(NotImplementedError):
                 self.gauge_transformer.as_sweep(input_circuit, N=1)
             return

diff --git a/cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test_utils_test.py b/cirq-core/cirq/transformers/gauge_compiling/gauge_compiling_test_utils_test.py
@@ -26,7 +26,15 @@ def _unitary_(self) -> np.ndarray:
         return self.unitary
 
 
+class ExampleSweepGate(cirq.testing.TwoQubitGate):
+    unitary = cirq.unitary(cirq.CZ)
+
+    def _unitary_(self) -> np.ndarray:
+        return self.unitary
+
+
 _EXAMPLE_TARGET = ExampleGate()
+_EXAMPLE_SWEEP_TARGET = ExampleSweepGate()
 
 _GOOD_TRANSFORMER = GaugeTransformer(
     target=_EXAMPLE_TARGET,
@@ -40,6 +48,22 @@ def _unitary_(self) -> np.ndarray:
     ),
 )
 
+_TRANSFORMER_WITH_SWEEP = GaugeTransformer(
+    target=_EXAMPLE_SWEEP_TARGET,
+    gauge_selector=GaugeSelector(
+        gauges=[
+            ConstantGauge(
+                two_qubit_gate=_EXAMPLE_SWEEP_TARGET,
+                pre_q0=cirq.Z,
+                pre_q1=cirq.Z,
+                post_q0=cirq.Z,
+                post_q1=cirq.Z,
+                support_sweep=True,
+            )
+        ]
+    ),
+)
+
 
 class TestValidTransformer(GaugeTester):
     two_qubit_gate = _EXAMPLE_TARGET
@@ -50,3 +74,9 @@ class TestInvalidTransformer(GaugeTester):
     two_qubit_gate = _EXAMPLE_TARGET
     gauge_transformer = _BAD_TRANSFORMER
     must_fail = True
+
+
+class TestSweep(GaugeTester):
+    two_qubit_gate = _EXAMPLE_SWEEP_TARGET
+    gauge_transformer = _TRANSFORMER_WITH_SWEEP
+    sweep_must_pass = True