def mock_write_metadata(metadata, path):
      """Mocks metadata_io.write_metadata to fail the first time it is called by this test, thus forcing a retry which should succeed."""
      if not write_metadata_called_list:
        write_metadata_called_list.append(True)
        original_write_metadata(metadata, path)
        raise ArithmeticError('Some error')
      return original_write_metadata(metadata, path)
Mocks metadata_io.write_metadata to fail the first time it is called by this test, thus forcing a retry which should succeed.
tensorflow_transform/beam/tft_beam_io/beam_metadata_io_test.py