def test_bulk_write_error(self, mock_client):
# Configure the mock client such that it raises BulkWriteError
# on insertion of records.
mock_method = mock.Mock(side_effect=errors.BulkWriteError(None))
mock_client()['fake_db']['fake_record_type'].insert_many = mock_method
# Ensure that the error is handled gracefully without a crash.
store = mongodbstore.MongoDBStore()
store.write({'record_type': 'foo', 'i': 0})
store.done()
def test_write_with_flush(self, mock_client):
# Set buffer size of 3.
store = mongodbstore.MongoDBStore(buffer_size=3)
# Then insert 3 records.
for i in range(3):
record = {'record_type': 'foo', 'i': i}
store.write(record)
# Verify that buffer is being flushed via client.
mock_method = mock_client()['fake_db']['fake_record_type'].insert_many
mock_method.assert_called_once_with(mock.ANY, ordered=mock.ANY)
def test_write_without_flush(self, mock_client):
# Set buffer size of 3.
store = mongodbstore.MongoDBStore(buffer_size=3)
# But insert only 2 records.
for i in range(2):
record = {'record_type': 'foo', 'i': i}
store.write(record)
# Verify that buffer is not being flushed via client.
mock_method = mock_client()['fake_db']['fake_record_type'].insert_many
mock_method.assert_not_called()
def test_done(self, mock_client):
# Write only one record to store with default buffer size.
store = mongodbstore.MongoDBStore()
store.write({'record_type': 'foo', 'i': 0})
# Call store's done(), so that it flushes the pending buffer.
store.done()
# Verify that the pending buffer is flushed via client.
mock_method = mock_client()['fake_db']['fake_record_type'].insert_many
mock_method.assert_called_once_with(mock.ANY, ordered=mock.ANY)
# Verify that client is being closed.
mock_method = mock_client().close
mock_method.assert_called_once_with()