def test_flush(self, inlet, outlet):
buffer = Buffer(count_threshold=100, time_threshold=10)
payload = [1, 2, 3, 4]
records = [Record(payload=p) for p in payload]
link = Link(inlet,
outlet,
interval=1,
processors=buffer,
copy_records=False)
inlet._pull = pull_mock(records[:2])
link.transfer()
outlet._push.assert_called_with([], mock.ANY) # no records yet
inlet._pull = pull_mock(records[2:])
link.transfer()
outlet._push.assert_called_with([], mock.ANY) # no records yet
buffer.flush = True
inlet._pull = pull_mock([])
link.transfer()
outlet._push.assert_called_with(
records, mock.ANY) # all records should be flushed
def test_flush(self):
buffer = Buffer(count_threshold=100, time_threshold=10)
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
buffer.flush = False
rvA = buffer(records[:6])
buffer.flush = True
rvB = buffer(records[6:])
self.assertEqual(rvA, [], 'Should not contain any records yet')
self.assertEqual(rvB, records, 'Should contain all records')
def test_count_controller(self):
buffer = Buffer(count_threshold=10)
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
result_false = buffer._count_controller(records[:6])
result_true = buffer._count_controller(records)
self.assertFalse(
result_false,
'Count controller should return False with 6 records')
self.assertTrue(result_true,
'Count controller should return True with 11 records')
def test_on_reset(self):
custom_controller = MagicMock()
on_reset = MagicMock()
buffer = Buffer(custom_controllers=custom_controller,
on_reset=on_reset)
payload = [1, 2]
records = [Record(payload=p) for p in payload]
buffer(records)
custom_controller.assert_called_with(records)
self.assertEqual(on_reset.call_count, 1)
buffer.reset()
self.assertEqual(on_reset.call_count, 2)
def test_time_controller(self):
buffer = Buffer(time_threshold=0.01)
payload = [1, 2, 3]
records = [Record(payload=p) for p in payload]
result_false = buffer._time_controller(records)
time.sleep(0.01)
result_true = buffer._time_controller(records)
self.assertFalse(
result_false,
'Time controller should return False after less than 0.01 seconds')
self.assertTrue(
result_true,
'Time controller should return True after sleeping for 0.01 seconds'
)
def test_execute_conjoint_true(self):
custom_controller = MagicMock(side_effect=lambda x: True)
buffer = Buffer(count_threshold=1,
time_threshold=1,
custom_controllers=custom_controller,
conjugate_controllers=True)
buffer._count_controller = MagicMock(side_effect=lambda x: True)
buffer._time_controller = MagicMock(side_effect=lambda x: True)
payload = [1, 2, 3]
records = [Record(payload=p) for p in payload]
result = buffer._execute(records)
buffer._count_controller.assert_called_with(records)
buffer._time_controller.assert_called_with(records)
custom_controller.assert_called_with(records)
self.assertEqual(result, records)
def test_buffer_time(self, inlet, outlet):
buffer = Buffer(time_threshold=0.02)
payload = [1, 2, 3, 4]
records = [Record(payload=p) for p in payload]
link = Link(inlet,
outlet,
interval=1,
processors=buffer,
copy_records=False)
inlet._pull = pull_mock(records[:2])
link.transfer()
outlet._push.assert_called_with(
[], mock.ANY) # not enough time have passed
inlet._pull = pull_mock(records[2:])
link.transfer()
outlet._push.assert_called_with(
[], mock.ANY) # not enough time have passed
time.sleep(0.02)
inlet._pull = pull_mock([])
link.transfer()
outlet._push.assert_called_with(
records, mock.ANY) # all records should be returned here
def test_execute_count(self):
buffer = Buffer(count_threshold=10)
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
rvA = buffer(records[:6])
rvB = buffer(records[6:])
self.assertEqual(rvA, [], 'Should not contain any records yet')
self.assertEqual(rvB, records, 'Should contain all records')
def test_flush_after_shutdown(self, inlet, outlet):
buffer = Buffer(count_threshold=100, time_threshold=10)
counter_dict = {'counter': 0, 'records': []}
link = Link(inlet,
outlet,
interval=0.01,
processors=buffer,
copy_records=False)
planner = SchedulePlanner(link, refresh_interval=0.01)
async def pull_coro(_):
counter_dict['counter'] += 1
record = Record(payload=counter_dict['counter'])
counter_dict['records'].append(record)
return [record]
mock_pull = MagicMock(side_effect=pull_coro)
inlet._pull = mock_pull
th = Thread(target=planner.start, daemon=True)
th.start()
time.sleep(0.1)
planner.shutdown()
th.join()
calls = outlet._push.call_args_list
for c in calls:
self.assertEqual(c(), [],
'Should only contain empty record lists.')
self.assertEqual(buffer.records, counter_dict['records'],
'All records should be stored in the buffer')
planner.force_transfer()
self.assertEqual(outlet._push.call_args(), [],
'Should return empty record list')
buffer.flush = True
planner.force_transfer()
self.assertEqual(outlet._push.call_args[0][0], counter_dict['records'],
'Should return all records')
def test_execute_custom(self):
# checks if number 11 is in the any of the records' payload
custom_controller = lambda x: len(
list(filter(lambda v: v.payload == 11, x)))
buffer = Buffer(custom_controllers=custom_controller)
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
rvA = buffer(records[:6])
rvB = buffer(records[6:])
self.assertEqual(rvA, [], 'Should not contain any records yet')
self.assertEqual(rvB, records, 'Should contain all records')
def test_reset(self):
buffer = Buffer()
buffer.records = [1, 2, 3]
buffer.time_start = time.time()
buffer.reset()
self.assertEqual(buffer.records, [],
'Buffer.records should be reset to []')
self.assertIsNone(buffer.time_start,
'Buffer.time_start should be reset to None')
def test_execute_count_before_time(self):
buffer = Buffer(count_threshold=10, time_threshold=1)
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
rv = []
for r in records:
rv = buffer([r])
if rv != []:
break
self.assertEqual(rv, records, 'Should contain all records')
self.assertEqual(buffer.records, [],
'Buffer.records should be reset to []')
self.assertIsNone(buffer.time_start,
'Buffer.time_start should be reset to None')
def test_custom_exception(self):
custom_controller = MagicMock(
side_effect=DummyException('Custom controller exception!'))
payload = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
records = [Record(payload=p) for p in payload]
buffer = Buffer(count_threshold=10,
custom_controllers=custom_controller)
with self.assertLogs(logging.getLogger('databay.Buffer'),
level='DEBUG') as cm:
rvA = buffer(records[:6])
rvB = buffer(records[6:])
self.assertTrue(
'Custom controller exception!' in ';'.join(cm.output))
self.assertEqual(rvA, [], 'Should not contain any records yet')
self.assertEqual(rvB, records, 'Should contain all records')
def test_execute_disjoint(self):
custom_controller = MagicMock(side_effect=lambda x: False)
buffer = Buffer(count_threshold=1,
time_threshold=1,
custom_controllers=custom_controller)
buffer._count_controller = MagicMock(side_effect=lambda x: False)
buffer._time_controller = MagicMock(side_effect=lambda x: False)
payload = [1, 2, 3]
records = [Record(payload=p) for p in payload]
buffer._execute(records)
buffer._count_controller.assert_called_with(records)
buffer._time_controller.assert_called_with(records)
custom_controller.assert_called_with(records)
def test_buffer_count(self, inlet, outlet):
buffer = Buffer(count_threshold=3)
payload = [1, 2, 3, 4]
records = [Record(payload=p) for p in payload]
link = Link(inlet,
outlet,
interval=1,
processors=buffer,
copy_records=False)
inlet._pull = pull_mock(records[:2])
link.transfer()
outlet._push.assert_called_with(
[], mock.ANY) # after first call we shouldn't have any records
inlet._pull = pull_mock(records[2:])
link.transfer()
outlet._push.assert_called_with(
records, mock.ANY) # all records should be returned here
def test_get_controllers(self):
custom_controller = lambda x: True
buffer = Buffer()
self.assertEqual(buffer.get_controllers(), [],
'Should not have any controllers')
buffer = Buffer(time_threshold=1)
self.assertEqual(buffer.get_controllers(), [buffer._time_controller],
'Should have only time controller')
buffer = Buffer(count_threshold=1)
self.assertEqual(buffer.get_controllers(), [buffer._count_controller],
'Should have only count controller')
buffer = Buffer(custom_controllers=[custom_controller])
self.assertEqual(buffer.get_controllers(), [custom_controller],
'Should have only custom controller')
buffer = Buffer(time_threshold=1,
count_threshold=1,
custom_controllers=[custom_controller])
self.assertEqual(buffer.get_controllers(), [
custom_controller, buffer._count_controller,
buffer._time_controller
], 'Should have all controllers')
