def mock_copy_tree_to_unique_temp_dir(source, base_temp_dir_path):
      """Mocks transform_fn_io._copy_tree to fail the first time it is called by this test, thus forcing a retry which should succeed."""
      global _COPY_TREE_TO_UNIQUE_TEMP_DIR_CALLED
      if not _COPY_TREE_TO_UNIQUE_TEMP_DIR_CALLED:
        _COPY_TREE_TO_UNIQUE_TEMP_DIR_CALLED = True
        original_copy_tree_to_unique_temp_dir(source, base_temp_dir_path)
        raise ArithmeticError('Some error')
      return original_copy_tree_to_unique_temp_dir(source, base_temp_dir_path)
Mocks transform_fn_io._copy_tree to fail the first time it is called by this test, thus forcing a retry which should succeed.
tensorflow_transform/beam/tft_beam_io/transform_fn_io_test.py