async def test_memory_checkpoint(self):
# first consumer
checkpointer = MemoryCheckPointer(name="test")
consumer_a = Consumer(stream_name=None,
checkpointer=checkpointer,
max_shard_consumers=1)
self.patch_consumer_fetch(consumer_a)
consumer_a.shards = [{"ShardId": "test-1"}, {"ShardId": "test-2"}]
await consumer_a.fetch()
shards = [s["ShardId"] for s in consumer_a.shards if s.get("stats")]
# Expect only one shard assigned as max = 1
self.assertEqual(["test-1"], shards)
# second consumer (note: max_shard_consumers needs to be 2 as uses checkpointer to get allocated shards)
consumer_b = Consumer(stream_name=None,
checkpointer=checkpointer,
max_shard_consumers=2)
self.patch_consumer_fetch(consumer_b)
consumer_b.shards = [{"ShardId": "test-1"}, {"ShardId": "test-2"}]
await consumer_b.fetch()
shards = [s["ShardId"] for s in consumer_b.shards if s.get("stats")]
# Expect only one shard assigned as max = 1
self.assertEqual(["test-2"], shards)
async def test_memory_checkpoint(self, *args):
# *args pass through mock
# first consumer
checkpointer = MemoryCheckPointer(name="test")
consumer_a = Consumer(
stream_name=None,
checkpointer=checkpointer,
max_shard_consumers=1,
endpoint_url=ENDPOINT_URL,
)
consumer_a.get_shard_iterator.return_value = True
consumer_a.shards = [{
"ShardId": "test-1",
"SequenceNumberRange": {}
}, {
"ShardId": "test-2",
"SequenceNumberRange": {}
}]
consumer_a.stream_status = consumer_a.ACTIVE
consumer_a.shards_status = consumer_a.INITIALIZE
await consumer_a.sync_shards()
shards = [s["ShardId"] for s in consumer_a.shards if s.get("stats")]
# Expect only one shard assigned as max = 1
self.assertEqual(["test-1"], shards)
# second consumer (note: max_shard_consumers needs to be 2 as uses checkpointer to get allocated shards)
consumer_b = Consumer(
stream_name=None,
checkpointer=checkpointer,
max_shard_consumers=2,
endpoint_url=ENDPOINT_URL,
)
consumer_b.get_shard_iterator.return_value = True
consumer_b.shards = [{
"ShardId": "test-1",
"SequenceNumberRange": {}
}, {
"ShardId": "test-2",
"SequenceNumberRange": {}
}]
consumer_b.stream_status = consumer_b.ACTIVE
consumer_b.shards_status = consumer_b.INITIALIZE
await consumer_b.sync_shards()
shards = [s["ShardId"] for s in consumer_b.shards if s.get("stats")]
# Expect only one shard assigned as max = 2
self.assertEqual(["test-1", "test-2"], shards)