def test_balance_ownership_balanced(ownerships, partitions, expected_result):
ownership_ref = {
'ownership_active0': {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "0",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_0",
"last_modified_time": time.time()
},
'ownership_active1': {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "1",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_1",
"last_modified_time": time.time()
},
'ownership_self_owned': {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "1",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": TEST_OWNER,
"last_modified_time": time.time()
},
'ownership_expired': {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "2",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "owner_1",
"last_modified_time": time.time() - 100000
},
'ownership_released': {
"fully_qualified_namespace": TEST_NAMESPACE,
"partition_id": "3",
"eventhub_name": TEST_EVENTHUB,
"consumer_group": TEST_CONSUMER_GROUP,
"owner_id": "",
"last_modified_time": time.time()
}
}
class MockEventHubClient(object):
eventhub_name = TEST_EVENTHUB
def __init__(self):
self._address = _Address(hostname=TEST_NAMESPACE,
path=MockEventHubClient.eventhub_name)
def get_partition_ids(self):
return ["0", "1"]
mock_client = MockEventHubClient()
current_ownerships = [ownership_ref[o] for o in ownerships]
om = OwnershipManager(mock_client, TEST_CONSUMER_GROUP, TEST_OWNER, None,
10, LoadBalancingStrategy.BALANCED, None)
to_claim_ownership = om._balance_ownership(current_ownerships, partitions)
assert len(to_claim_ownership) == expected_result
async def test_ownership_manager_release_partition():
class MockEventHubClient(object):
eventhub_name = "test_eh_name"
def __init__(self):
self._address = _Address(hostname="test",
path=MockEventHubClient.eventhub_name)
async def get_partition_ids(self):
return ["0", "1"]
class MockCheckpointStore(InMemoryCheckpointStore):
released = None
async def claim_ownership(self, ownsership):
self.released = ownsership
checkpoint_store = MockCheckpointStore()
ownership_manager = OwnershipManager(MockEventHubClient(), "$Default",
"owner", checkpoint_store, 10.0, "0")
ownership_manager.cached_parition_ids = ["0", "1"]
ownership_manager.owned_partitions = []
await ownership_manager.release_ownership("1")
assert checkpoint_store.released is None
ownership_manager.owned_partitions = [{
"partition_id":
"0",
"owner_id":
"foo",
"last_modified_time":
time.time() + 31
}]
await ownership_manager.release_ownership("0")
assert checkpoint_store.released is None
ownership_manager.owned_partitions = [{
"partition_id": "0",
"owner_id": "",
"last_modified_time": time.time()
}]
await ownership_manager.release_ownership("0")
assert checkpoint_store.released is None
ownership_manager.owned_partitions = [{
"partition_id": "0",
"owner_id": "foo",
"last_modified_time": time.time()
}]
await ownership_manager.release_ownership("0")
assert checkpoint_store.released is None
ownership_manager.owned_partitions = [{
"partition_id": "0",
"owner_id": "owner",
"last_modified_time": time.time()
}]
await ownership_manager.release_ownership("0")
assert checkpoint_store.released[0]["owner_id"] == ""