cleanup _assert_equivalent_op_tree (#3909)

Adds cirq.testing.assert_equivalent_op_tree.

This method was called from three other testing packages as a private method, this change moves it into its proper place: cirq.testing.

cirq/contrib/noise_models/noise_models_test.py

@@ -15,7 +15,7 @@
 import cirq
 import cirq.contrib.noise_models as ccn
 from cirq import ops
-from cirq.devices.noise_model_test import _assert_equivalent_op_tree
+from cirq.testing import assert_equivalent_op_tree
 
 
 def test_depol_noise():
@@ -83,7 +83,7 @@ def test_readout_noise_after_moment():
             cirq.measure(qubits[2], key='q2'),
         ]
     )
-    _assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
+    assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
 
 
 # Composes depolarization, damping, and readout noise (in that order).
@@ -140,7 +140,7 @@ def test_decay_noise_after_moment():
             cirq.measure(qubits[2], key='q2'),
         ]
     )
-    _assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
+    assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
 
 
 # Test the aggregate noise models.
@@ -188,7 +188,7 @@ def test_aggregate_readout_noise_after_moment():
             cirq.measure(qubits[2], key='q2'),
         ]
     )
-    _assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
+    assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
 
 
 def test_aggregate_decay_noise_after_moment():
@@ -238,4 +238,4 @@ def test_aggregate_decay_noise_after_moment():
             cirq.measure(qubits[2], key='q2'),
         ]
     )
-    _assert_equivalent_op_tree(true_noisy_program, noisy_circuit)
+    assert_equivalent_op_tree(true_noisy_program, noisy_circuit)

cirq/devices/noise_model_test.py

@@ -20,18 +20,13 @@
 import cirq
 from cirq import ops
 from cirq.devices.noise_model import validate_all_measurements
+from cirq.testing import assert_equivalent_op_tree
 
 
-def _assert_equivalent_op_tree(x: cirq.OP_TREE, y: cirq.OP_TREE):
-    a = list(cirq.flatten_op_tree(x))
-    b = list(cirq.flatten_op_tree(y))
-    assert a == b
-
-
-def _assert_equivalent_op_tree_sequence(x: Sequence[cirq.OP_TREE], y: Sequence[cirq.OP_TREE]):
+def assert_equivalent_op_tree_sequence(x: Sequence[cirq.OP_TREE], y: Sequence[cirq.OP_TREE]):
     assert len(x) == len(y)
     for a, b in zip(x, y):
-        _assert_equivalent_op_tree(a, b)
+        assert_equivalent_op_tree(a, b)
 
 
 def test_requires_one_override():
@@ -58,11 +53,11 @@ def noisy_moments(self, moments, system_qubits):
             return result
 
     a = NoiseModelWithNoisyMomentListMethod()
-    _assert_equivalent_op_tree(a.noisy_operation(cirq.H(q)), cirq.X(q).with_tags(ops.VirtualTag()))
-    _assert_equivalent_op_tree(
+    assert_equivalent_op_tree(a.noisy_operation(cirq.H(q)), cirq.X(q).with_tags(ops.VirtualTag()))
+    assert_equivalent_op_tree(
         a.noisy_moment(cirq.Moment([cirq.H(q)]), [q]), cirq.X(q).with_tags(ops.VirtualTag())
     )
-    _assert_equivalent_op_tree_sequence(
+    assert_equivalent_op_tree_sequence(
         a.noisy_moments([cirq.Moment(), cirq.Moment([cirq.H(q)])], [q]),
         [[], cirq.X(q).with_tags(ops.VirtualTag())],
     )
@@ -72,11 +67,11 @@ def noisy_moment(self, moment, system_qubits):
             return [y.with_tags(ops.VirtualTag()) for y in cirq.Y.on_each(*moment.qubits)]
 
     b = NoiseModelWithNoisyMomentMethod()
-    _assert_equivalent_op_tree(b.noisy_operation(cirq.H(q)), cirq.Y(q).with_tags(ops.VirtualTag()))
-    _assert_equivalent_op_tree(
+    assert_equivalent_op_tree(b.noisy_operation(cirq.H(q)), cirq.Y(q).with_tags(ops.VirtualTag()))
+    assert_equivalent_op_tree(
         b.noisy_moment(cirq.Moment([cirq.H(q)]), [q]), cirq.Y(q).with_tags(ops.VirtualTag())
     )
-    _assert_equivalent_op_tree_sequence(
+    assert_equivalent_op_tree_sequence(
         b.noisy_moments([cirq.Moment(), cirq.Moment([cirq.H(q)])], [q]),
         [[], cirq.Y(q).with_tags(ops.VirtualTag())],
     )
@@ -86,11 +81,11 @@ def noisy_operation(self, operation: 'cirq.Operation'):
             return cirq.Z(operation.qubits[0]).with_tags(ops.VirtualTag())
 
     c = NoiseModelWithNoisyOperationMethod()
-    _assert_equivalent_op_tree(c.noisy_operation(cirq.H(q)), cirq.Z(q).with_tags(ops.VirtualTag()))
-    _assert_equivalent_op_tree(
+    assert_equivalent_op_tree(c.noisy_operation(cirq.H(q)), cirq.Z(q).with_tags(ops.VirtualTag()))
+    assert_equivalent_op_tree(
         c.noisy_moment(cirq.Moment([cirq.H(q)]), [q]), cirq.Z(q).with_tags(ops.VirtualTag())
     )
-    _assert_equivalent_op_tree_sequence(
+    assert_equivalent_op_tree_sequence(
         c.noisy_moments([cirq.Moment(), cirq.Moment([cirq.H(q)])], [q]),
         [[], cirq.Z(q).with_tags(ops.VirtualTag())],
     )
@@ -155,8 +150,8 @@ def test_noise_composition():
     merge(actual_zs)
     merge(actual_sz)
     merge(expected_circuit)
-    _assert_equivalent_op_tree(actual_zs, actual_sz)
-    _assert_equivalent_op_tree(actual_zs, expected_circuit)
+    assert_equivalent_op_tree(actual_zs, actual_sz)
+    assert_equivalent_op_tree(actual_zs, expected_circuit)
 
 
 def test_constant_qubit_noise_repr():

cirq/google/experimental/noise_models/noise_models_test.py

@@ -19,7 +19,7 @@
 
 import cirq
 from cirq import ops
-from cirq.devices.noise_model_test import _assert_equivalent_op_tree
+from cirq.testing import assert_equivalent_op_tree
 from cirq.google.api import v2
 from cirq.google.experimental.noise_models import (
     simple_noise_from_calibration_metrics,
@@ -168,7 +168,7 @@ def test_per_qubit_depol_noise_from_data():
             ]
         ),
     )
-    _assert_equivalent_op_tree(expected_program, noisy_circuit)
+    assert_equivalent_op_tree(expected_program, noisy_circuit)
 
 
 def test_per_qubit_readout_error_from_data():
@@ -214,7 +214,7 @@ def test_per_qubit_readout_error_from_data():
             ]
         ),
     )
-    _assert_equivalent_op_tree(expected_program, noisy_circuit)
+    assert_equivalent_op_tree(expected_program, noisy_circuit)
 
 
 def test_per_qubit_readout_decay_from_data():
@@ -255,7 +255,7 @@ def test_per_qubit_readout_decay_from_data():
             ]
         ),
     )
-    _assert_equivalent_op_tree(expected_program, noisy_circuit)
+    assert_equivalent_op_tree(expected_program, noisy_circuit)
 
 
 def test_per_qubit_combined_noise_from_data():
@@ -319,4 +319,4 @@ def test_per_qubit_combined_noise_from_data():
             ]
         ),
     )
-    _assert_equivalent_op_tree(expected_program, noisy_circuit)
+    assert_equivalent_op_tree(expected_program, noisy_circuit)

cirq/testing/__init__.py

@@ -112,3 +112,5 @@
 from cirq.testing.deprecation import (
     assert_deprecated,
 )
+
+from cirq.testing.op_tree import assert_equivalent_op_tree

cirq/testing/op_tree.py

@@ -0,0 +1,32 @@
+# Copyright 2020 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from cirq import ops
+
+
+def assert_equivalent_op_tree(x: ops.OP_TREE, y: ops.OP_TREE):
+    """Ensures that the two OP_TREEs are equivalent.
+
+    Args:
+        x: OP_TREE one
+        y: OP_TREE two
+    Returns:
+        None
+    Raises:
+         AssertionError if x != y
+    """
+
+    a = list(ops.flatten_op_tree(x))
+    b = list(ops.flatten_op_tree(y))
+    assert a == b

cirq/testing/op_tree_test.py

@@ -0,0 +1,29 @@
+# Copyright 2020 The Cirq Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+
+import cirq
+from cirq.testing import assert_equivalent_op_tree
+
+
+def test_assert_equivalent_op_tree():
+    assert_equivalent_op_tree([], [])
+    a = cirq.NamedQubit("a")
+    assert_equivalent_op_tree([cirq.X(a)], [cirq.X(a)])
+
+    assert_equivalent_op_tree(cirq.Circuit([cirq.X(a)]), [cirq.X(a)])
+    assert_equivalent_op_tree(cirq.Circuit([cirq.X(a)], cirq.Moment()), [cirq.X(a)])
+
+    with pytest.raises(AssertionError):
+        assert_equivalent_op_tree([cirq.X(a)], [])