async def test_claim_existing_ownership():
ownership = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0",
}
checkpoint_store = InMemoryCheckpointStore()
await checkpoint_store.claim_ownership([ownership])
listed_ownership = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
existing_ownership = listed_ownership[0]
etag_existing = existing_ownership["etag"]
last_modified_exising = existing_ownership["last_modified_time"]
time.sleep(0.01) # so time is changed
copied_existing_ownership = existing_ownership.copy()
copied_existing_ownership["owner_id"] = "owner_1"
await checkpoint_store.claim_ownership([copied_existing_ownership])
listed_ownership2 = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
assert len(listed_ownership2) == 1
assert listed_ownership2[0] == copied_existing_ownership
assert listed_ownership2[0]["owner_id"] == "owner_1"
assert listed_ownership2[0]["etag"] != etag_existing
assert listed_ownership2[0]["last_modified_time"] != last_modified_exising
async def test_receive_load_balancing_async(connstr_senders):
connection_str, senders = connstr_senders
cs = InMemoryCheckpointStore()
client1 = EventHubConsumerClient.from_connection_string(
connection_str,
consumer_group='$default',
checkpoint_store=cs,
load_balancing_interval=1)
client2 = EventHubConsumerClient.from_connection_string(
connection_str,
consumer_group='$default',
checkpoint_store=cs,
load_balancing_interval=1)
async def on_event(partition_context, event):
pass
async with client1, client2:
task1 = asyncio.ensure_future(
client1.receive(on_event, starting_position="-1"))
await asyncio.sleep(3.3)
task2 = asyncio.ensure_future(
client2.receive(on_event, starting_position="-1"))
await asyncio.sleep(10)
assert len(client1._event_processors[("$default",
ALL_PARTITIONS)]._tasks) == 1
assert len(client2._event_processors[("$default",
ALL_PARTITIONS)]._tasks) == 1
await task1
await task2
async def test_partition_processor_process_eventhub_consumer_error():
async def event_handler(partition_context, event):
pass
async def error_handler(partition_context, error):
error_handler.error = error
async def partition_close_handler(partition_context, reason):
partition_close_handler.reason = reason
class MockEventHubClient(object):
eventhub_name = "test_eh_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
def _create_consumer(self, consumer_group, partition_id,
event_position, on_event_received, **kwargs):
return MockEventhubConsumer(on_event_received=on_event_received,
**kwargs)
async def get_partition_ids(self):
return ["0", "1"]
class MockEventhubConsumer(object):
def __init__(self, **kwargs):
self.stop = False
self._on_event_received = kwargs.get("on_event_received")
async def receive(self, *args, **kwargs):
raise EventHubError("Mock EventHubConsumer EventHubError")
async def close(self):
pass
eventhub_client = MockEventHubClient()
checkpoint_store = InMemoryCheckpointStore()
event_processor = EventProcessor(
eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=error_handler,
partition_initialize_handler=None,
partition_close_handler=partition_close_handler,
load_balancing_interval=1.3)
task = asyncio.ensure_future(event_processor.start())
await asyncio.sleep(5)
await event_processor.stop()
await task
assert isinstance(error_handler.error, EventHubError)
assert partition_close_handler.reason == CloseReason.OWNERSHIP_LOST
async def test_claim_new_ownership():
ownership = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0"
}
checkpoint_store = InMemoryCheckpointStore()
claimed_ownership = await checkpoint_store.claim_ownership([ownership])
assert len(claimed_ownership) == 1
listed_ownership = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
assert listed_ownership[0] == ownership
assert ownership["last_modified_time"] is not None
assert ownership["etag"] is not None
async def test_update_checkpoint():
checkpoint = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"offset": "0",
"sequencenumber": 0
}
checkpoint_store = InMemoryCheckpointStore()
await checkpoint_store.update_checkpoint(checkpoint)
listed_checkpoint = await checkpoint_store.list_checkpoints(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
assert listed_checkpoint[0] == checkpoint
assert listed_checkpoint[0]["offset"] == "0"
assert listed_checkpoint[0]["sequencenumber"] == 0
async def test_claim_two_partitions():
ownership_0 = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0"
}
ownership_1 = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "1",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0"
}
checkpoint_store = InMemoryCheckpointStore()
await checkpoint_store.claim_ownership([ownership_0, ownership_1])
listed_ownership0 = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
assert set(x["partition_id"] for x in listed_ownership0) == {"0", "1"}
async def test_partition_processor_process_events_error(connstr_senders):
async def event_handler(partition_context, event):
if partition_context.partition_id == "1":
raise RuntimeError("processing events error")
else:
pass
async def error_handler(partition_context, error):
if partition_context.partition_id == "1":
error_handler.error = error
else:
raise RuntimeError(
"There shouldn't be an error for partition other than 1")
async def partition_close_handler(partition_context, reason):
if partition_context.partition_id == "1":
assert reason == CloseReason.OWNERSHIP_LOST
else:
assert reason == CloseReason.SHUTDOWN
connection_str, senders = connstr_senders
for sender in senders:
sender.send(EventData("EventProcessor Test"))
eventhub_client = EventHubConsumerClient.from_connection_string(
connection_str, consumer_group='$default')
checkpoint_store = InMemoryCheckpointStore()
event_processor = EventProcessor(
eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=error_handler,
initial_event_position="-1",
partition_initialize_handler=None,
partition_close_handler=partition_close_handler,
load_balancing_interval=1)
task = asyncio.ensure_future(event_processor.start())
await asyncio.sleep(10)
await event_processor.stop()
# task.cancel()
await eventhub_client.close()
assert isinstance(error_handler.error, RuntimeError)
async def test_claim_two_ownerships():
ownership_0 = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0"
}
ownership_1 = {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP+"1",
"owner_id": "owner_1"
}
checkpoint_store = InMemoryCheckpointStore()
claimed_ownership = await checkpoint_store.claim_ownership([ownership_0, ownership_1])
assert len(claimed_ownership) == 2
listed_ownership0 = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP)
listed_ownership1 = await checkpoint_store.list_ownership(TEST_NAMESPACE, TEST_EVENTHUB, TEST_CONSUMER_GROUP+"1")
assert listed_ownership0[0]["consumer_group"] == TEST_CONSUMER_GROUP
assert listed_ownership1[0]["consumer_group"] == TEST_CONSUMER_GROUP+"1"
async def test_partition_processor_process_error_close_error():
async def partition_initialize_handler(partition_context):
partition_initialize_handler.called = True
raise RuntimeError("initialize error")
async def event_handler(partition_context, event):
event_handler.called = True
raise RuntimeError("process_events error")
async def error_handler(partition_context, error):
assert isinstance(error, RuntimeError)
error_handler.called = True
raise RuntimeError("process_error error")
async def partition_close_handler(partition_context, reason):
assert reason == CloseReason.OWNERSHIP_LOST
partition_close_handler.called = True
raise RuntimeError("close error")
class MockEventHubClient(object):
eventhub_name = "test_eh_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
def _create_consumer(self, consumer_group, partition_id,
event_position, on_event_received, **kwargs):
return MockEventhubConsumer(on_event_received=on_event_received,
**kwargs)
async def get_partition_ids(self):
return ["0", "1"]
class MockEventhubConsumer(object):
def __init__(self, **kwargs):
self.stop = False
self._on_event_received = kwargs.get("on_event_received")
async def receive(self, *args, **kwargs):
await asyncio.sleep(0.1)
await self._on_event_received(EventData("mock events"))
async def close(self):
pass
class MockOwnershipManager(OwnershipManager):
called = False
async def release_ownership(self, partition_id):
self.called = True
eventhub_client = MockEventHubClient()
checkpoint_store = InMemoryCheckpointStore()
ownership_manager = MockOwnershipManager(eventhub_client, "$Default",
"owner", checkpoint_store, 10.0,
"0")
event_processor = EventProcessor(
eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=error_handler,
partition_initialize_handler=partition_initialize_handler,
partition_close_handler=partition_close_handler,
load_balancing_interval=1.3)
event_processor._ownership_manager = ownership_manager
task = asyncio.ensure_future(event_processor.start())
await asyncio.sleep(5)
await event_processor.stop()
await task
assert partition_initialize_handler.called
assert event_handler.called
assert error_handler.called
assert ownership_manager.called
assert partition_close_handler.called
async def test_loadbalancer_balance():
class MockEventHubClient(object):
eventhub_name = "test_eventhub_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
def _create_consumer(self, consumer_group, partition_id,
event_position, on_event_received, **kwargs):
consumer = MockEventhubConsumer(
on_event_received=on_event_received, **kwargs)
return consumer
async def get_partition_ids(self):
return ["0", "1"]
async def close(self):
pass
class MockEventhubConsumer(object):
def __init__(self, **kwargs):
self.stop = False
self._on_event_received = kwargs.get("on_event_received")
async def receive(self, *args, **kwargs):
await asyncio.sleep(0.1)
await self._on_event_received(EventData("mock events"))
async def close(self):
pass
eventhub_client = MockEventHubClient()
checkpoint_store = InMemoryCheckpointStore()
tasks = []
event_processor1 = EventProcessor(eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=None,
partition_initialize_handler=None,
partition_close_handler=None,
load_balancing_interval=1.3)
tasks.append(asyncio.ensure_future(event_processor1.start()))
await asyncio.sleep(3)
assert len(
event_processor1._tasks) == 2 # event_processor1 claims two partitions
event_processor2 = EventProcessor(eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=None,
partition_initialize_handler=None,
partition_close_handler=None,
load_balancing_interval=1.3)
tasks.append(asyncio.ensure_future(event_processor2.start()))
await asyncio.sleep(3)
assert len(event_processor1._tasks
) == 1 # two event processors balance. So each has 1 task
assert len(event_processor2._tasks) == 1
event_processor3 = EventProcessor(eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=None,
partition_initialize_handler=None,
partition_close_handler=None,
load_balancing_interval=1.3)
tasks.append(asyncio.ensure_future(event_processor3.start()))
await asyncio.sleep(3)
assert len(event_processor3._tasks) == 0
await event_processor3.stop()
await event_processor1.stop()
await asyncio.sleep(3)
assert len(
event_processor2._tasks
) == 2 # event_procesor2 takes another one after event_processor1 stops
await event_processor2.stop()
await eventhub_client.close()
async def test_partition_processor_process_events_error():
class MockEventHubClient(object):
eventhub_name = "test_eventhub_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
def _create_consumer(self, consumer_group, partition_id,
event_position, on_event_received, **kwargs):
consumer = MockEventhubConsumer(
on_event_received=on_event_received, **kwargs)
return consumer
async def get_partition_ids(self):
return ["0", "1"]
async def close(self):
pass
class MockEventhubConsumer(object):
def __init__(self, **kwargs):
self.stop = False
self._on_event_received = kwargs.get("on_event_received")
async def receive(self, *args, **kwargs):
await asyncio.sleep(0.1)
await self._on_event_received(EventData("mock events"))
async def close(self):
pass
async def event_handler(partition_context, event):
if partition_context.partition_id == "1":
raise RuntimeError("processing events error")
else:
pass
async def error_handler(partition_context, error):
if partition_context.partition_id == "1":
error_handler.error = error
else:
raise RuntimeError(
"There shouldn't be an error for partition other than 1")
eventhub_client = MockEventHubClient()
checkpoint_store = InMemoryCheckpointStore()
event_processor = EventProcessor(eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=error_handler,
initial_event_position="-1",
partition_initialize_handler=None,
load_balancing_interval=1.3)
task = asyncio.ensure_future(event_processor.start())
await asyncio.sleep(10)
await event_processor.stop()
await task
await asyncio.sleep(1)
await eventhub_client.close()
assert isinstance(error_handler.error, RuntimeError)
async def test_partition_processor():
class MockEventHubClient(object):
eventhub_name = "test_eventhub_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
def _create_consumer(self, consumer_group, partition_id,
event_position, on_event_received, **kwargs):
consumer = MockEventhubConsumer(
on_event_received=on_event_received, **kwargs)
return consumer
async def get_partition_ids(self):
return ["0", "1"]
async def close(self):
pass
class MockEventhubConsumer(object):
def __init__(self, **kwargs):
self.stop = False
self._on_event_received = kwargs.get("on_event_received")
async def receive(self, *args, **kwargs):
await asyncio.sleep(0.1)
await self._on_event_received(EventData("mock events"))
async def close(self):
pass
lock = asyncio.Lock()
event_map = {}
checkpoint = None
close_reason = None
error = None
async def partition_initialize_handler(partition_context):
partition_initialize_handler.partition_context = partition_context
async def event_handler(partition_context, event):
async with lock:
if event:
nonlocal checkpoint, event_map
event_map[partition_context.partition_id] = event_map.get(
partition_context.partition_id, 0) + 1
offset, sn = event.offset, event.sequence_number
checkpoint = (offset, sn)
await partition_context.update_checkpoint(event)
async def partition_close_handler(partition_context, reason):
assert partition_context and reason
nonlocal close_reason
close_reason = reason
async def error_handler(partition_context, err):
assert partition_context and err
nonlocal error
error = err
eventhub_client = MockEventHubClient()
checkpoint_store = InMemoryCheckpointStore()
event_processor = EventProcessor(
eventhub_client=eventhub_client,
consumer_group='$default',
checkpoint_store=checkpoint_store,
event_handler=event_handler,
error_handler=error_handler,
partition_initialize_handler=partition_initialize_handler,
partition_close_handler=partition_close_handler,
load_balancing_interval=1.3)
task = asyncio.ensure_future(event_processor.start())
await asyncio.sleep(2)
assert len(event_processor._tasks) == 2
await event_processor.stop()
await task
await eventhub_client.close()
assert event_map['0'] >= 1 and event_map['1'] >= 1
assert checkpoint is not None
assert close_reason == CloseReason.SHUTDOWN
assert error is None
assert partition_initialize_handler.partition_context
