def test_adaptor_with_ownership(self, take_ownership):
mock_writer = mock.MagicMock()
adaptor = io.RawProtoWriterAdaptor(mock_writer,
take_ownership=take_ownership)
# Write out protos to our adaptor.
with adaptor as enter_return_value:
# Make sure that __enter__ returns the adaptor itself.
self.assertIs(adaptor, enter_return_value)
adaptor.write(self.proto1)
adaptor.write(self.proto2)
if take_ownership:
# If we took ownership, mock_writer __enter__ and __exit__ should have
# been called.
mock_writer.__enter__.assert_called_once_with()
test_utils.assert_called_once_workaround(mock_writer.__exit__)
else:
# If not, they shouldn't have been called.
test_utils.assert_not_called_workaround(mock_writer.__enter__)
test_utils.assert_not_called_workaround(mock_writer.__exit__)
self.assertEqual(
mock_writer.write.call_args_list,
[mock.call(r.SerializeToString()) for r in self.protos])
def __init__(self, options):
self._writers = {k: None for k in ['candidates', 'examples', 'gvcfs']}
if options.candidates_filename:
self._add_writer('candidates',
io_utils.RawProtoWriterAdaptor(
io_utils.make_tfrecord_writer(
options.candidates_filename)))
if options.examples_filename:
self._add_writer('examples',
io_utils.RawProtoWriterAdaptor(
io_utils.make_tfrecord_writer(
options.examples_filename)))
if options.gvcf_filename:
self._add_writer('gvcfs',
io_utils.RawProtoWriterAdaptor(
io_utils.make_tfrecord_writer(
options.gvcf_filename)))