def test_record_batches_rec_too_large(self, failure_mock):
"""FirehoseClient - Record Batches, Record Too Large"""
records = [{'key': 'test' * 1000 * 1000}]
result = list(FirehoseClient._record_batches(records))
assert_equal(result, [[]])
failure_mock.assert_called_with(1)
def test_record_batches_max_batch_count(self):
"""FirehoseClient - Record Batches, Max Batch Count"""
records = self._sample_raw_records(count=501)
result = list(FirehoseClient._record_batches(records))
assert_equal(len(result), 2)
assert_equal(len(result[0]), 500)
assert_equal(len(result[1]), 1)
def test_record_batches(self):
"""FirehoseClient - Record Batches"""
records = self._sample_raw_records()
expected_result = [['{"key_0":"value_0"}\n', '{"key_1":"value_1"}\n']]
result = list(FirehoseClient._record_batches(records))
assert_equal(result, expected_result)
def test_record_batches_max_batch_size(self):
"""FirehoseClient - Record Batches, Max Batch Size"""
records = [{'key_{}'.format(i): 'test' * 100000} for i in range(10)]
result = list(FirehoseClient._record_batches(records))
assert_equal(len(result), 2)
assert_equal(len(result[0]), 9)
assert_equal(len(result[1]), 1)
batch_size_01 = sum(len(rec) for rec in result[0])
batch_size_02 = sum(len(rec) for rec in result[1])
assert_equal(batch_size_01 < FirehoseClient.MAX_BATCH_SIZE, True)
assert_equal(batch_size_02 < FirehoseClient.MAX_BATCH_SIZE, True)
assert_equal(
batch_size_01 + batch_size_02 > FirehoseClient.MAX_BATCH_SIZE,
True)
