Add arolla.testing.override_operator context manager.

Thanks Mihai for the original implementation.

PiperOrigin-RevId: 871419074
Change-Id: Ia253d1f58c3b9cb8982dab9a7a81299d93c46bb4

Author: timofey-stepanov

py/koladata/type_checking/type_checking_test.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import contextlib
 from absl.testing import absltest
 from absl.testing import parameterized
 from arolla import arolla
@@ -33,21 +32,6 @@
 UuPerson = kd.schema.uu_schema(age=kd.INT32, name=kd.STRING)
 
 
-@contextlib.contextmanager
-def override_operator(operator_name: str, new_operator: arolla.abc.Operator):
-  old_operator = arolla.abc.decay_registered_operator(operator_name)
-  arolla.abc.register_operator(
-      operator_name, new_operator, if_present='unsafe_override'
-  )
-
-  try:
-    yield
-  finally:
-    arolla.abc.register_operator(
-        operator_name, old_operator, if_present='unsafe_override'
-    )
-
-
 ERROR_CASES = (
     (
         '_primitive',
@@ -720,7 +704,9 @@ def get_repr(x):  # pylint: disable=unused-argument
           '`kd.schema.get_repr` should not have been called',
       )
 
-    with override_operator('kd.schema.get_repr', get_repr):
+    with arolla.testing.override_registered_operator(
+        'kd.schema.get_repr', get_repr
+    ):
       _ = fn(ds([1]))  # Does not raise.
 
   @parameterized.named_parameters(('boolean', kd.BOOLEAN), ('none', None))