#------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. #-------------------------------------------------------------------------- import logging import os import pytest import asyncio import sys import time import uamqp from uamqp import address, types, utils, authentication, MessageBodyType from uamqp.message import DataBody, ValueBody, SequenceBody def get_logger(level): uamqp_logger = logging.getLogger("uamqp") if not uamqp_logger.handlers: handler = logging.StreamHandler(stream=sys.stdout) handler.setFormatter(logging.Formatter('%(asctime)s %(name)-12s %(levelname)-8s %(message)s')) uamqp_logger.addHandler(handler) uamqp_logger.setLevel(level) return uamqp_logger log = get_logger(logging.INFO) def on_message_received(message): annotations = message.annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) return message def send_single_message(live_eventhub_config, partition, msg_content): target = "amqps://{}/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], partition ) uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAuth.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) uamqp.send_message(target, msg_content, auth=sas_auth, debug=False) def send_multiple_message(live_eventhub_config, msg_count): def data_generator(): for i in range(msg_count): msg_content = "Hello world {}".format(i).encode('utf-8') yield msg_content uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAuth.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) target = "amqps://{}/{}/Partitions/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['partition']) send_client = uamqp.SendClient(target, auth=sas_auth, debug=False) message_batch = uamqp.message.BatchMessage(data_generator()) send_client.queue_message(message_batch) results = send_client.send_all_messages(close_on_done=False) assert not [m for m in results if m == uamqp.constants.MessageState.SendFailed] send_client.close() @pytest.mark.asyncio async def test_event_hubs_callback_async_receive(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) receive_client = uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=1000, prefetch=10) log.info("Created client, receiving...") await receive_client.receive_messages_async(on_message_received) log.info("Finished receiving") @pytest.mark.asyncio async def test_event_hubs_callback_async_receive_no_shutdown_after_timeout(live_eventhub_config): received_cnt = {'cnt': 0} def on_message_received_internal(message): annotations = message.annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) log.info(str(message)) message.accept() received_cnt['cnt'] += 1 uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) source = address.Source(source) source.set_filter(b"amqp.annotation.x-opt-offset > '@latest'") receive_client = uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=3000, prefetch=10, shutdown_after_timeout=False) log.info("Created client, receiving...") try: await receive_client.open_async() while not await receive_client.client_ready_async(): await asyncio.sleep(0.05) await asyncio.sleep(1) # sleep for 1s await receive_client._connection.work_async() # do a single connection iteration to see if there're incoming transfers # make sure no messages are received assert not receive_client._was_message_received assert receive_client._received_messages.empty() send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') await receive_client.receive_messages_async(on_message_received_internal) message_handler_before = receive_client.message_handler assert received_cnt['cnt'] == 1 send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') await receive_client.receive_messages_async(on_message_received_internal) message_handler_after = receive_client.message_handler assert message_handler_before == message_handler_after assert received_cnt['cnt'] == 2 log.info("Finished receiving") finally: await receive_client.close_async() @pytest.mark.asyncio async def test_event_hubs_filter_receive_async(live_eventhub_config): plain_auth = authentication.SASLPlain( live_eventhub_config['hostname'], live_eventhub_config['key_name'], live_eventhub_config['access_key']) source_url = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) source = address.Source(source_url) source.set_filter(b"amqp.annotation.x-opt-enqueuedtimeutc > 1518731960545") receive_client = uamqp.ReceiveClientAsync(source, auth=plain_auth, timeout=5000) await receive_client.receive_messages_async(on_message_received) @pytest.mark.asyncio async def test_event_hubs_iter_receive_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) receive_client = uamqp.ReceiveClientAsync(source, debug=False, auth=sas_auth, timeout=3000, prefetch=10) count = 0 message_generator = receive_client.receive_messages_iter_async() async for message in message_generator: log.info("No. {} : {}".format(message.annotations.get(b'x-opt-sequence-number'), message)) count += 1 if count >= 10: log.info("Got {} messages. Breaking.".format(count)) message.accept() break count = 0 async for message in message_generator: count += 1 if count >= 10: log.info("Got {} more messages. Shutting down.".format(count)) message.accept() break await receive_client.close_async() @pytest.mark.asyncio async def test_event_hubs_iter_receive_no_shutdown_after_timeout_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) source = address.Source(source) source.set_filter(b"amqp.annotation.x-opt-offset > '@latest'") receive_client = uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=5000, debug=False, shutdown_after_timeout=False) count = 0 try: await receive_client.open_async() while not await receive_client.client_ready_async(): await asyncio.sleep(0.05) await asyncio.sleep(1) # sleep for 1s await receive_client._connection.work_async() # do a single connection iteration to see if there're incoming transfers # make sure no messages are received assert not receive_client._was_message_received assert receive_client._received_messages.empty() gen = receive_client.receive_messages_iter_async() send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') async for message in gen: log.info(message.annotations.get(b'x-opt-sequence-number')) log.info(str(message)) count += 1 assert count == 1 count = 0 message_handler_before = receive_client.message_handler send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') gen = receive_client.receive_messages_iter_async() async for message in gen: log.info(message.annotations.get(b'x-opt-sequence-number')) log.info(str(message)) count += 1 assert count == 1 message_handler_after = receive_client.message_handler assert message_handler_before == message_handler_after finally: await receive_client.close_async() @pytest.mark.asyncio async def test_event_hubs_batch_receive_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) async with uamqp.ReceiveClientAsync(source, debug=False, auth=sas_auth, timeout=3000, prefetch=10) as receive_client: message_batch = await receive_client.receive_message_batch_async(10) log.info("got batch: {}".format(len(message_batch))) for message in message_batch: annotations = message.annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) next_batch = await receive_client.receive_message_batch_async(10) log.info("got another batch: {}".format(len(next_batch))) for message in next_batch: annotations = message.annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) next_batch = await receive_client.receive_message_batch_async(10) log.info("got another batch: {}".format(len(next_batch))) for message in next_batch: annotations = message.annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) @pytest.mark.asyncio async def test_event_hubs_batch_receive_async_no_shutdown_after_timeout_sync(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) source = address.Source(source) source.set_filter(b"amqp.annotation.x-opt-offset > '@latest'") async with uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=3000, prefetch=10, shutdown_after_timeout=False) as receive_client: received_cnt = 0 received_cnt += len(await receive_client.receive_message_batch_async(10)) assert received_cnt == 0 message_handler_before = receive_client.message_handler send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') received_cnt += len(await receive_client.receive_message_batch_async(10)) assert received_cnt == 1 received_cnt += len(await receive_client.receive_message_batch_async(10)) assert received_cnt == 1 send_single_message(live_eventhub_config, live_eventhub_config['partition'], 'message') received_cnt += len(await receive_client.receive_message_batch_async(10)) message_handler_after = receive_client.message_handler assert message_handler_before == message_handler_after assert received_cnt == 2 @pytest.mark.asyncio async def test_event_hubs_client_web_socket_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key'], transport_type=uamqp.TransportType.AmqpOverWebsocket) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) async with uamqp.ReceiveClientAsync(source, auth=sas_auth, debug=False, timeout=5000, prefetch=50) as receive_client: receive_client.receive_message_batch(max_batch_size=10) @pytest.mark.asyncio async def test_event_hubs_receive_with_runtime_metric_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) receiver_runtime_metric_symbol = b'com.microsoft:enable-receiver-runtime-metric' symbol_array = [types.AMQPSymbol(receiver_runtime_metric_symbol)] desired_capabilities = utils.data_factory(types.AMQPArray(symbol_array)) async with uamqp.ReceiveClientAsync(source, debug=False, auth=sas_auth, timeout=1000, prefetch=10, desired_capabilities=desired_capabilities) as receive_client: message_batch = await receive_client.receive_message_batch_async(10) log.info("got batch: {}".format(len(message_batch))) for message in message_batch: annotations = message.annotations delivery_annotations = message.delivery_annotations log.info("Sequence Number: {}".format(annotations.get(b'x-opt-sequence-number'))) assert b'last_enqueued_sequence_number' in delivery_annotations assert b'last_enqueued_offset' in delivery_annotations assert b'last_enqueued_time_utc' in delivery_annotations assert b'runtime_info_retrieval_time_utc' in delivery_annotations @pytest.mark.asyncio async def test_event_hubs_shared_connection_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group']) async with uamqp.ConnectionAsync(live_eventhub_config['hostname'], sas_auth, debug=False) as conn: partition_0 = uamqp.ReceiveClientAsync(source + "0", debug=False, auth=sas_auth, timeout=3000, prefetch=10) partition_1 = uamqp.ReceiveClientAsync(source + "1", debug=False, auth=sas_auth, timeout=3000, prefetch=10) await partition_0.open_async(connection=conn) await partition_1.open_async(connection=conn) tasks = [ partition_0.receive_message_batch_async(1), partition_1.receive_message_batch_async(1) ] try: messages = await asyncio.gather(*tasks) assert len(messages[0]) == 1 and len(messages[1]) == 1 except: raise finally: await partition_0.close_async() await partition_1.close_async() async def receive_ten(partition, receiver): messages = [] count = 0 while count < 10: print("Receiving {} on partition {}".format(count, partition)) batch = await receiver.receive_message_batch_async(1) print("Received {} messages on partition {}".format(len(batch), partition)) messages.extend(batch) count += 1 print("Finished receiving on partition {}".format(partition)) return messages @pytest.mark.asyncio async def test_event_hubs_multiple_receiver_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth_a = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) sas_auth_b = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group']) partition_0 = uamqp.ReceiveClientAsync(source + "0", debug=False, auth=sas_auth_a, timeout=3000, prefetch=10) partition_1 = uamqp.ReceiveClientAsync(source + "1", debug=False, auth=sas_auth_b, timeout=3000, prefetch=10) try: await partition_0.open_async() await partition_1.open_async() tasks = [ receive_ten("0", partition_0), receive_ten("1", partition_1) ] messages = await asyncio.gather(*tasks) assert len(messages) == 2 assert len(messages[0]) >= 10 assert len(messages[1]) >= 10 print(messages) finally: await partition_0.close_async() await partition_1.close_async() @pytest.mark.asyncio async def test_event_hubs_dynamic_issue_link_credit_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) msg_sent_cnt = 200 send_multiple_message(live_eventhub_config, msg_sent_cnt) def message_received_callback(message): message_received_callback.received_msg_cnt += 1 message_received_callback.received_msg_cnt = 0 async with uamqp.ReceiveClientAsync(source, debug=True, auth=sas_auth, prefetch=1) as receive_client: receive_client._message_received_callback = message_received_callback while not await receive_client.client_ready_async(): await asyncio.sleep(0.05) await asyncio.sleep(1) # sleep for 1s await receive_client._connection.work_async() # do a single connection iteration to see if there're incoming transfers # make sure no messages are received assert not receive_client._was_message_received assert receive_client._received_messages.empty() await receive_client.message_handler.reset_link_credit_async(msg_sent_cnt) now = start = time.time() wait_time = 5 while now - start <= wait_time: await receive_client._connection.work_async() now = time.time() assert message_received_callback.received_msg_cnt == msg_sent_cnt log.info("Finished receiving") @pytest.mark.asyncio async def test_event_hubs_not_receive_events_during_connection_establishment_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], live_eventhub_config['partition']) receive_client = uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=1000, debug=False, prefetch=10) try: await receive_client.open_async() while not await receive_client.client_ready_async(): await asyncio.sleep(0.05) await asyncio.sleep(1) # sleep for 1s await receive_client._connection.work_async() # do a single connection iteration to see if there're incoming transfers # make sure no messages are received assert not receive_client._was_message_received assert receive_client._received_messages.empty() messages_0 = await receive_client.receive_message_batch_async() assert len(messages_0) > 0 finally: await receive_client.close_async() @pytest.mark.asyncio async def event_hubs_send_different_amqp_body_type_async(live_eventhub_config): uri = "sb://{}/{}".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) target = "amqps://{}/{}/Partitions/0".format(live_eventhub_config['hostname'], live_eventhub_config['event_hub']) send_client = uamqp.SendClientAsync(target, auth=sas_auth, debug=False) data_body_1 = [b'data1', b'data2'] data_body_message_1 = uamqp.message.Message(body=data_body_1) data_body_message_1.application_properties = {'body_type': 'data_body_1'} send_client.queue_message(data_body_message_1) data_body_2 = b'data1' data_body_message_2 = uamqp.message.Message(body=data_body_2, body_type=MessageBodyType.Data) data_body_message_2.application_properties = {'body_type': 'data_body_2'} send_client.queue_message(data_body_message_2) value_body_1 = [b'data1', -1.23, True, {b'key': b'value'}, [1, False, 1.23, b'4']] value_body_message_1 = uamqp.message.Message(body=value_body_1) value_body_message_1.application_properties = {'body_type': 'value_body_1'} send_client.queue_message(value_body_message_1) value_body_2 = {b'key1': {b'sub_key': b'sub_value'}, b'key2': b'value', 3: -1.23} value_body_message_2 = uamqp.message.Message(body=value_body_2, body_type=MessageBodyType.Value) value_body_message_2.application_properties = {'body_type': 'value_body_2'} send_client.queue_message(value_body_message_2) sequence_body_1 = [b'data1', -1.23, True, {b'key': b'value'}, [b'a', 1.23, True]] sequence_body_message_1 = uamqp.message.Message(body=sequence_body_1, body_type=MessageBodyType.Sequence) sequence_body_message_1.application_properties = {'body_type': 'sequence_body_1'} send_client.queue_message(sequence_body_message_1) sequence_body_2 = [[1, 2, 3], [b'aa', b'bb', b'cc'], [True, False, True], [{b'key1': b'value'}, {b'key2': 123}]] sequence_body_message_2 = uamqp.message.Message(body=sequence_body_2, body_type=MessageBodyType.Sequence) sequence_body_message_2.application_properties = {'body_type': 'sequence_body_2'} send_client.queue_message(sequence_body_message_2) results = await send_client.send_all_messages_async(close_on_done=False) assert not [m for m in results if m == uamqp.constants.MessageState.SendFailed] sas_auth = authentication.SASTokenAsync.from_shared_access_key( uri, live_eventhub_config['key_name'], live_eventhub_config['access_key']) source = "amqps://{}/{}/ConsumerGroups/{}/Partitions/{}".format( live_eventhub_config['hostname'], live_eventhub_config['event_hub'], live_eventhub_config['consumer_group'], 0) result_dic = {} receive_client = uamqp.ReceiveClientAsync(source, auth=sas_auth, timeout=5000, debug=False, prefetch=10) gen = receive_client.receive_messages_iter_async() async for message in gen: if message.application_properties and message.application_properties.get(b'body_type'): if message.application_properties.get(b'body_type') == b'data_body_1': check_list = [data for data in message.get_data()] assert isinstance(message._body, DataBody) assert check_list == data_body_1 result_dic['data_body_1'] = 1 elif message.application_properties.get(b'body_type') == b'data_body_2': check_list = [data for data in message.get_data()] assert isinstance(message._body, DataBody) assert check_list == [data_body_2] result_dic['data_body_2'] = 1 elif message.application_properties.get(b'body_type') == b'value_body_1': assert message.get_data() == value_body_1 assert isinstance(message._body, ValueBody) result_dic['value_body_1'] = 1 elif message.application_properties.get(b'body_type') == b'value_body_2': assert message.get_data() == value_body_2 assert isinstance(message._body, ValueBody) result_dic['value_body_2'] = 1 elif message.application_properties.get(b'body_type') == b'sequence_body_1': check_list = [data for data in message.get_data()] assert check_list == [sequence_body_1] assert isinstance(message._body, SequenceBody) result_dic['sequence_body_1'] = 1 elif message.application_properties.get(b'body_type') == b'sequence_body_2': check_list = [data for data in message.get_data()] assert check_list == sequence_body_2 assert isinstance(message._body, SequenceBody) result_dic['sequence_body_2'] = 1 log.info(message.annotations.get(b'x-opt-sequence-number')) log.info(str(message)) await send_client.close_async() await receive_client.close_async() assert len(results) == 6 if __name__ == '__main__': config = {} config['hostname'] = os.environ['EVENT_HUB_HOSTNAME'] config['event_hub'] = os.environ['EVENT_HUB_NAME'] config['key_name'] = os.environ['EVENT_HUB_SAS_POLICY'] config['access_key'] = os.environ['EVENT_HUB_SAS_KEY'] config['consumer_group'] = "$Default" config['partition'] = "0" loop = asyncio.get_event_loop() loop.run_until_complete(event_hubs_send_different_amqp_body_type_async(config))