def test_invalid_json_error():
def callback(_, __):
"""
empty callback func
:param _:
:param __:
:return:
"""
print('done')
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
def on_gap(_, __):
"""
should not detect a gap
:param _:
:param __:
:return:
"""
raise Exception('GapDetectError')
sub.on('gap', on_gap)
msg1 = create_msg()
msg3 = create_msg(3, 2)
sub.handle_message(msg1)
error = InvalidJsonError(stream_id, 'invalid json', 'test error msg',
create_msg(2, 1))
sub.handle_error(error)
sub.handle_message(msg3)
def test_no_gap():
count = 0
def callback(_, __):
"""
callback function to test whether has received msg
:param _:
:param __:
:return:
"""
nonlocal count
count += 1
msg1 = create_msg()
msg2 = create_msg(2, 1)
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
def on_gap(_, __):
"""
gap should not be detected
:param _:
:param __:
:return:
"""
raise Exception('GapDetectError')
sub.on('gap', on_gap)
sub.handle_message(msg1)
sub.handle_message(msg2)
assert count == 2
def test_gap():
count = 0
def callback(_, __):
"""
callback function to test whether has received msg
:param _:
:param __:
:return:
"""
nonlocal count
count += 1
msg1 = create_msg()
msg4 = create_msg(4, 3)
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
def on_gap(from_, to_):
"""
test whether gap is detected
:param from_:
:param to_:
:return:
"""
assert from_ == 2
assert to_ == 3
print('gap detected')
sub.on('gap', on_gap)
sub.handle_message(msg1)
sub.handle_message(msg4)
assert count == 1
def test_emit_event():
count = 0
sub = Subscription(stream_id, stream_partition, 'api_key',
lambda: print('done'))
def test():
nonlocal count
count += 1
sub.on(EventConstant.SUBSCRIBED, test)
sub.set_state(EventConstant.SUBSCRIBED)
assert count == 1
def test_handle_error():
err = Exception('Test error')
def callback(_, __):
"""
should not be called
:param _:
:param __:
:return:
"""
raise Exception('handleErrorFailed')
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
def test(msg):
assert isinstance(msg, Exception)
assert str(msg) == 'Test error'
sub.on('error', test)
sub.handle_error(err)
def test_bye():
count = 0
def callback(_, __):
"""
callback function to test whether has received msg
:param _:
:param __:
:return:
"""
nonlocal count
count += 1
bye_msg = create_msg(1, None, {'BYE_KEY': True})
sub = Subscription(bye_msg.stream_id, bye_msg.stream_partition, 'api_key',
callback)
def test():
assert count == 1
sub.on('done', test)
sub.handle_message(bye_msg)
def subscribe(self, stream, callback, legacy_option=None):
"""
subscribe to stream with given id
:param stream: object or dict contains stream_id and stream_partition
:param callback: callback function when subscribed
:param legacy_option: backward compatibility
:return: subscription
"""
if hasattr(stream, 'stream_id'):
stream_id = stream.stream_id
elif isinstance(stream, dict) and ('stream_id' in stream.keys()
or 'stream' in stream.keys()):
stream_id = stream.get('stream', None) or stream.get(
'stream_id', None)
elif isinstance(stream, str):
stream_id = stream
else:
raise ValueError(
'subscribe: stream_id and stream_partition should be given. Given :%s'
% (type(stream)))
if isinstance(legacy_option, Option):
opt = copy.copy(legacy_option)
opt.stream_id = stream_id
else:
opt = Option()
sub = Subscription(stream_id, opt.stream_partition or 0,
self.option.api_key, callback, opt)
def gap_handler(from_, to_):
"""
handler when subscription detect a gap
:param from_:
:param to_:
:return: None
"""
if not sub.resending:
self.__request_resend(sub,
Option(resend_from=from_, resend_to=to_))
sub.on(EventConstant.GAP, gap_handler)
def done_handler():
"""
handler when subscription is done
:return: None
"""
logger.debug('done event for sub %s ' % sub.sub_id)
self.unsubscribe(sub)
sub.on(EventConstant.DONE, done_handler)
self.__add_subscription(sub)
if self.connection.state == EventConstant.CONNECTED:
self.__resend_and_subscribe(sub)
elif self.option.auto_connect is True:
try:
self.connect()
except Exception as e:
import traceback
traceback.print_exc()
raise ConnectionErr(e)
return sub