async def test_producer_and_consumer_consume_with_bytes(self):
class ByteSerializer(Serializer):
def serialize(self, msg):
result = str.encode(msg)
return result
def deserialize(self, data):
return data
class ByteProcessor(Processor, NetstringAggregator, ByteSerializer):
pass
processor = ByteProcessor()
async with Producer(
stream_name=self.stream_name, endpoint_url=ENDPOINT_URL, processor=processor
) as producer:
for x in range(0, 2):
await producer.put(f"{x}")
await producer.flush()
results = []
checkpointer = MemoryCheckPointer(name="test")
async with Consumer(
stream_name=self.stream_name,
endpoint_url=ENDPOINT_URL,
processor=processor,
checkpointer=checkpointer,
) as consumer:
async for item in consumer:
results.append(item)
await checkpointer.checkpoint(
shard_id=consumer.shards[0]["ShardId"], sequence="seq"
)
async for item in consumer:
results.append(item)
self.assertEquals(len(results), 2)
await checkpointer.close()
self.assertEquals(len(checkpointer.get_all_checkpoints()), 1)
async def test_consumer_checkpoint(self):
checkpointer = MemoryCheckPointer(name="test")
results = []
async with Producer(
stream_name=self.STREAM_NAME_SINGLE_SHARD,
processor=StringProcessor(),
) as producer:
async with Consumer(
stream_name=self.STREAM_NAME_SINGLE_SHARD,
checkpointer=checkpointer,
processor=StringProcessor(),
iterator_type="LATEST",
) as consumer:
# Manually start
await consumer.start_consumer()
await producer.put("test")
await producer.flush()
for i in range(3):
async for item in consumer:
results.append(item)
checkpoints = checkpointer.get_all_checkpoints()
# Expect 1 as only 1 shard
self.assertEquals(1, len(checkpoints))
self.assertIsNotNone(checkpoints[list(
checkpoints.keys())[0]]["sequence"])
self.assertListEqual(results, ["test"])
async def test_memory_checkpoint(self):
# first consumer
checkpointer = MemoryCheckPointer(name="test")
consumer_a = Consumer(
stream_name=None,
checkpointer=checkpointer,
max_shard_consumers=1,
endpoint_url=ENDPOINT_URL,
)
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,
endpoint_url=ENDPOINT_URL,
)
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_producer_and_consumer_consume_with_checkpointer_and_latest(
self):
async with Producer(stream_name=self.stream_name,
endpoint_url=ENDPOINT_URL) as producer:
await producer.put("test.A")
results = []
checkpointer = MemoryCheckPointer(name="test")
async with Consumer(
stream_name=self.stream_name,
endpoint_url=ENDPOINT_URL,
checkpointer=checkpointer,
iterator_type="LATEST",
) as consumer:
async for item in consumer:
results.append(item)
# Expect none as LATEST
self.assertEquals([], results)
checkpoints = checkpointer.get_all_checkpoints()
# Expect 1 as only 1 shard
self.assertEquals(1, len(checkpoints))
# none as no records yet (using LATEST)
self.assertIsNone(checkpoints[list(
checkpoints.keys())[0]]["sequence"])
results = []
log.info("checkpointer checkpoints: {}".format(checkpoints))
log.info("Starting consumer again..")
async with Consumer(
stream_name=self.stream_name,
endpoint_url=ENDPOINT_URL,
checkpointer=checkpointer,
iterator_type="LATEST",
sleep_time_no_records=0.5,
) as consumer:
# Manually start
await consumer.start_consumer()
await producer.put("test.B")
await producer.flush()
log.info("waiting..")
await asyncio.sleep(1)
log.info("about to consume..")
async for item in consumer:
results.append(item)
self.assertEquals(["test.B"], results)
checkpoints = checkpointer.get_all_checkpoints()
log.info("checkpointer checkpoints: {}".format(checkpoints))
# expect not None as has processed records
self.assertIsNotNone(checkpoints[list(
checkpoints.keys())[0]]["sequence"])
# now add some records
for i in range(0, 10):
await producer.put("test.{}".format(i))
await producer.flush()
await asyncio.sleep(1)
results = []
async with Consumer(
stream_name=self.stream_name,
endpoint_url=ENDPOINT_URL,
checkpointer=checkpointer,
iterator_type="LATEST",
sleep_time_no_records=0.5,
) as consumer:
async for item in consumer:
results.append(item)
# Expect results as checkpointer resumed from prior sequence
self.assertEquals(10, len(results))
