def test__get_output_prop_kwargs_content_type_not_set(self, source_type):
mock = Mock(spec=BaseCollector,
type=source_type,
get_output_message_id=Mock(return_value=SAMPLE_MESSAGE_ID))
created_timestamp = 1234
with patch('time.time', return_value=created_timestamp) as time_mock:
# collectors handling the sources of "stream" type do not
# need the `content_type` to be set
if source_type == 'stream':
output_prop_kwargs = BaseCollector.get_output_prop_kwargs(
mock,
source=SAMPLE_SOURCE,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
arg_a=SAMPLE_ARG_A)
# assertions
time_mock.assert_called_once_with()
mock.get_output_message_id.assert_called_once_with(
source=SAMPLE_SOURCE,
created_timestamp=created_timestamp,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
arg_a=SAMPLE_ARG_A)
self.assertEqual(
output_prop_kwargs, {
'message_id': SAMPLE_MESSAGE_ID,
'type': source_type,
'timestamp': 1234,
'headers': {}
})
else:
with self.assertRaises(AttributeError):
BaseCollector.get_output_prop_kwargs(
mock,
source=SAMPLE_SOURCE,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
arg_a=SAMPLE_ARG_A)
def test__get_output_prop_kwargs(self, source_type):
mock = Mock(__class__=BaseCollector,
type=source_type,
content_type=SAMPLE_CONTENT_TYPE,
get_output_message_id=Mock(return_value=SAMPLE_MESSAGE_ID))
created_timestamp = 1234
with patch('time.time', return_value=created_timestamp) as time_mock:
# the call
output_prop_kwargs = BaseCollector.get_output_prop_kwargs(
mock,
source=SAMPLE_SOURCE,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
arg_a=SAMPLE_ARG_A)
# assertions
time_mock.assert_called_once_with()
mock.get_output_message_id.assert_called_once_with(
source=SAMPLE_SOURCE,
created_timestamp=created_timestamp,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
arg_a=SAMPLE_ARG_A)
# if the stream is of the type "source" - it does not
# add `content_type` to the properties
if source_type == 'stream':
self.assertEqual(output_prop_kwargs, {
'message_id': SAMPLE_MESSAGE_ID,
'type': source_type,
'timestamp': 1234,
'headers': {}})
else:
self.assertEqual(output_prop_kwargs, {
'message_id': SAMPLE_MESSAGE_ID,
'type': source_type,
'content_type': SAMPLE_CONTENT_TYPE,
'timestamp': 1234,
'headers': {}})
def test__get_source(self):
mock = Mock(__class__=BaseCollector,
config=dict(source='my_src_label'))
source = BaseCollector.get_source(mock,
'my_src_channel',
blablabla=sentinel.blablabla)
self.assertEqual(source, 'my_src_label.my_src_channel')
def test__get_output_components(self):
mock = Mock(
__class__=BaseCollector,
process_input_data=Mock(
return_value=dict(ccc=sentinel.ccc, dddd=sentinel.dddd)),
get_source_channel=Mock(return_value=SAMPLE_SOURCE_CHANNEL),
get_source=Mock(return_value=SAMPLE_SOURCE),
get_output_rk=Mock(return_value=SAMPLE_OUTPUT_RK),
get_output_data_body=Mock(return_value=SAMPLE_OUTPUT_DATA_BODY),
get_output_prop_kwargs=Mock(
return_value=SAMPLE_OUTPUT_PROP_KWARGS))
# the call
(output_rk, output_data_body,
output_prop_kwargs) = BaseCollector.get_output_components(
mock, a=SAMPLE_ARG_A, bb=SAMPLE_ARG_B)
# assertions
self.assertIs(output_rk, SAMPLE_OUTPUT_RK)
self.assertIs(output_data_body, SAMPLE_OUTPUT_DATA_BODY)
self.assertIs(output_prop_kwargs, SAMPLE_OUTPUT_PROP_KWARGS)
self.assertEqual(mock.mock_calls, [
call.process_input_data(a=SAMPLE_ARG_A, bb=SAMPLE_ARG_B),
call.get_source_channel(ccc=sentinel.ccc, dddd=sentinel.dddd),
call.get_source(source_channel=SAMPLE_SOURCE_CHANNEL,
ccc=sentinel.ccc,
dddd=sentinel.dddd),
call.get_output_rk(
source=SAMPLE_SOURCE, ccc=sentinel.ccc, dddd=sentinel.dddd),
call.get_output_data_body(
source=SAMPLE_SOURCE, ccc=sentinel.ccc, dddd=sentinel.dddd),
call.get_output_prop_kwargs(
source=SAMPLE_SOURCE,
output_data_body=SAMPLE_OUTPUT_DATA_BODY,
ccc=sentinel.ccc,
dddd=sentinel.dddd),
])
def test__get_output_rk__with__raw_format_version_tag(self):
mock = Mock(__class__=BaseCollector,
raw_format_version_tag='33333')
output_rk = BaseCollector.get_output_rk(mock,
'my_src_label.my_src_channel',
blablabla=sentinel.blablabla)
self.assertEqual(output_rk, 'my_src_label.my_src_channel.33333')
def test__process_input_data(self):
mock = Mock(__class__=BaseCollector)
processed_data = BaseCollector.process_input_data(mock,
a=SAMPLE_ARG_A,
bb=SAMPLE_ARG_B)
self.assertEqual(processed_data, dict(a=SAMPLE_ARG_A,
bb=SAMPLE_ARG_B))
def test__get_output_message_id(self):
source = 'my_src_label.my_src_channel'
created_timestamp = 1234
created_timestamp_str = '1234'
output_data_body = '1234'
mock = Mock(__class__=BaseCollector)
message_id = BaseCollector.get_output_message_id(
mock, source, created_timestamp, output_data_body)
### XXX CR: rather hardcode a few specific md5s instead of:
expected_message_id = hashlib.md5(source + '\0' +
created_timestamp_str + '\0' +
output_data_body).hexdigest()
self.assertEqual(message_id, expected_message_id)
def test__get_output_data_body(self):
mock = Mock(__class__=BaseCollector)
with self.assertRaises(NotImplementedError):
BaseCollector.get_output_data_body(mock,
SAMPLE_SOURCE,
blablabla=sentinel.blablabla)
def test__get_source_channel(self):
mock = Mock(__class__=BaseCollector)
with self.assertRaises(NotImplementedError):
BaseCollector.get_source_channel(mock)
def test__run_handling__not_interrupted(self):
mock = Mock(__class__=BaseCollector)
BaseCollector.run_handling(mock)
mock.run.assert_called_once_with()
self.assertEqual(mock.stop.mock_calls, [])
def test__run_handling__interrupted(self):
mock = Mock(__class__=BaseCollector,
run=Mock(side_effect=KeyboardInterrupt))
BaseCollector.run_handling(mock)
mock.run.assert_called_once_with()
mock.stop.assert_called_once_with()
def test__get_script_init_kwargs(self):
self.assertEqual(BaseCollector.get_script_init_kwargs(), {})