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_handle_messages():
msgs = list(
map(lambda x: create_msg(x, None
if x is None else x - 1), [1, 2, 3, 4, 5]))
received = []
def callback(content, received_msg):
"""
callback function when subscription received data from server
:param content:
:param received_msg:
:return:
"""
received.append(received_msg)
if len(received) == 5:
for i in range(5):
assert msgs[i] is received[i]
assert content == {}
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
for m in msgs:
sub.handle_message(m)
def test_get_original_resend_option():
def callback(_, __):
"""
empty callback func
:param _:
:param __:
:return:
"""
print('done')
sub = Subscription(stream_id, stream_partition, 'api_key', callback,
Option(resend_all=True))
assert sub.get_effective_resend_option().resend_all is True
def test_get_resend_option_after_received_data3():
def callback(_, __):
"""
empty callback func
:param _:
:param __:
:return:
"""
print('done')
sub = Subscription(stream_id, stream_partition, 'api_key', callback,
Option(resend_from_time=time.time()))
sub.handle_message(create_msg(10))
dic = sub.get_effective_resend_option()
assert dic.resend_from == 11
def test_get_resend_option_after_received_data4():
msg = create_msg()
def callback(_, __):
"""
empty callback func
:param _:
:param __:
:return:
"""
print('done')
sub = Subscription(stream_id, stream_partition, 'api_key', callback,
Option(resend_last=10))
sub.handle_message(msg)
dic = sub.get_effective_resend_option()
assert dic.resend_last == 10
def test_handle_message():
msg = create_msg()
def callback(content, receive_msg):
"""
callback function when subscription received data from server
:param content:
:param receive_msg:
:return:
"""
assert isinstance(content, type(msg.get_parsed_content()))
assert content == msg.get_parsed_content()
assert receive_msg is msg
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
sub.handle_message(msg)
def test_duplicate_handling():
msg = create_msg()
count = 0
def callback(_, __):
"""
callback function to test whether has received msg
:param _:
:param __:
:return:
"""
nonlocal count
count += 1
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
sub.handle_message(msg)
sub.handle_message(msg)
assert count == 1
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_event_handling_no_resent():
count = 0
msg = create_msg()
def test(_, __):
nonlocal count
count += 1
sub = Subscription(stream_id, stream_partition, 'api_key', test)
sub.set_resending(True)
sub.handle_message(msg)
assert count == 0
sub.emit(EventConstant.NO_RESEND)
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_message_when_resending_and_is_resend_true():
msg = create_msg()
count = 0
def callback(_, __):
"""
callback function to test the number of received msg
:param _:
:param __:
:return:
"""
nonlocal count
count += 1
sub = Subscription(stream_id, stream_partition, 'api_key', callback)
sub.set_resending(True)
sub.handle_message(msg, True)
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
def test_update_state():
sub = Subscription(stream_id, stream_partition, 'api_key',
lambda: print('done'))
sub.set_state(EventConstant.SUBSCRIBED)
assert sub.get_state() == EventConstant.SUBSCRIBED
