def test_ignored_message_type_is_ignored(self, KafkaConsumer):
self.mock_consumer(
KafkaConsumer,
value=b'json:{"message":{"code":"NY","name":"New York"},"version":1,"type":"us-state"}',
)
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.add_ignored_message_type("us-state")
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
# Even though message is valid, the serializer should never get called since message type is explicitly ignored.
self.assertEqual(FakeStateSerializer.call_count, 0)
self.assertTrue("state" not in self.serializers)
def test_multi_shard_consume(self):
# Send a bunch of messages to a bunch of shards
key = 1
value = b'json:{"message":{"code":"NY","name":"New York"},"version":1,"type":"us-state"}'
client = self.make_stream_with_record(str(key), value, shard_count=20)
for i in range(100):
key += 1
client.put_record(
StreamName=TOPIC_STATES, Data=value, PartitionKey=str(key)
)
# Test the values sent to our serializer match the message
test = {"i": 0}
def save(ser):
self.assertEqual(ser.validated_data["code"], "NY")
self.assertEqual(ser.validated_data["name"], "New York")
test["i"] += 1
FakeStateSerializer = self.mock_state_serializer(save)
# Consume messages. Log should have 101 messages in it now.
consumer = Consumer(TOPIC_STATES)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=2000)
self.assertEqual(FakeStateSerializer.call_count, 101)
self.assertEqual(test["i"], 101)
def test_invalid_message_ignored(self, KafkaConsumer):
self.mock_consumer(KafkaConsumer,
value=b'json:{"message":{"code":"NYC","name":"New York"},"version":1,"type":"us-state"}')
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(self.serializers['state'].save.call_count, 0)
def test_normal_consume(self, KafkaConsumer):
# Make a fake consumer to generate a message
fake_kafka_consumer = self.mock_consumer(
KafkaConsumer,
value=b'json:{"message":{"code":"NY","name":"New York"},"version":1,"type":"us-state"}',
max_calls=100,
)
# Test the values sent to our serializer match the message
def save(ser):
self.assertEqual(ser.validated_data["code"], "NY")
self.assertEqual(ser.validated_data["name"], "New York")
FakeStateSerializer = self.mock_state_serializer(save)
# Consume a message
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
# Test the expected mocks where called
KafkaConsumer.assert_called_once_with(
auto_offset_reset="earliest",
bootstrap_servers=["kafka:9092"],
consumer_timeout_ms=500,
enable_auto_commit=False,
)
fake_kafka_consumer.partitions_for_topic.assert_called_once_with(TOPIC_STATES)
fake_kafka_consumer.assign.assert_called_once_with(
[
TopicPartition(partition=0, topic=TOPIC_STATES),
TopicPartition(partition=1, topic=TOPIC_STATES),
]
)
self.assertEqual(KafkaConsumer.call_count, 1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(fake_kafka_consumer.__next__.call_count, 1)
self.assertEqual(self.serializers["state"].save.call_count, 1)
consumer.run(iter_limit=1)
self.assertEqual(KafkaConsumer.call_count, 1)
self.assertEqual(FakeStateSerializer.call_count, 2)
self.assertEqual(fake_kafka_consumer.__next__.call_count, 2)
self.assertEqual(self.serializers["state"].save.call_count, 1)
def test_unknown_version_ignored(self):
self.make_stream_with_record(
"NY",
b'json:{"message":{"code":"NY","name":"New York"},"version":2,"type":"us-state"}',
)
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 0)
def test_unknown_version_throws(self, KafkaConsumer):
self.mock_consumer(KafkaConsumer,
value=b'json:{"message":{"code":"NY","name":"New York"},"version":2,"type":"us-state"}')
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500, throw_errors=True)
consumer.register(FakeStateSerializer)
with self.assertRaises(UnknownMessageVersionError):
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 0)
def test_invalid_message_throws(self):
self.make_stream_with_record(
"NY",
b'json:{"message":{"code":"NYC","name":"New York"},"version":1,"type":"us-state"}',
)
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500, throw_errors=True)
consumer.register(FakeStateSerializer)
with self.assertRaises(ValidationError):
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(self.serializers["state"].save.call_count, 0)
def test_normal_consume(self):
self.make_stream_with_record(
"NY",
b'json:{"message":{"code":"NY","name":"New York"},"version":1,"type":"us-state"}',
)
# Test the values sent to our serializer match the message
def save(ser):
self.assertEqual(ser.validated_data["code"], "NY")
self.assertEqual(ser.validated_data["name"], "New York")
FakeStateSerializer = self.mock_state_serializer(save)
# Consume a message
consumer = Consumer(TOPIC_STATES)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=10)
self.assertEqual(self.serializers["state"].save.call_count, 1)
consumer.run(iter_limit=10)
self.assertEqual(self.serializers["state"].save.call_count, 1)
def test_roundtrip_state(self):
def save(ser):
self.assertEqual(ser.validated_data['code'], 'NY')
self.assertEqual(ser.validated_data['name'], 'New York')
FakeStateSerializer = self.mock_state_serializer(save)
producer = Producer(TOPIC_STATES, StateSerializer)
record = producer.send({'code': 'NY', 'name': 'New York'})
self.assertEqual(record.topic, 'us-states')
self.assertEqual(record.partition, 0)
self.assertTrue(record.offset >= 0)
# producer.client.flush()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=1000)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(self.serializers['state'].save.call_count, 1)
def test_roundtrip_state(self):
def save(ser):
self.assertEqual(ser.validated_data["code"], "NY")
self.assertEqual(ser.validated_data["name"], "New York")
FakeStateSerializer = self.mock_state_serializer(save)
producer = Producer(TOPIC_STATES, StateSerializer)
record = producer.send({"code": "NY", "name": "New York"})
self.assertEqual(record.topic, "us-states")
self.assertEqual(record.partition, 0)
self.assertTrue(record.offset >= 0)
# producer.client.flush()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=1000)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(self.serializers["state"].save.call_count, 1)
def test_missing_message_ignored(self):
self.make_stream_with_record("NY", b'json:{"version":1}')
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 0)
def build_person_consumer():
consumer = Consumer('test')
consumer.register(PrometheusSerializer)
return consumer
def sale_consumer():
from .serializers import AccountingSaleSerializer
consumer = Consumer(settings.ACCOUNTING_SALE_TOPIC)
consumer.register(AccountingSaleSerializer)
return consumer
def new_sale_consumer():
from .serializers import SaleSerializer
consumer = Consumer(settings.NEW_SALE_TOPIC)
consumer.register(SaleSerializer)
return consumer
def build_person_consumer():
consumer = Consumer(constants.TOPIC_PEOPLE)
consumer.register(serializers.PersonSerializer)
return consumer
def build_ride_consumer():
# Run using python manage.py run_kafka_consumer
consumer = Consumer('ride')
consumer.register(RideSerializer)
return consumer
def ticket_consumer():
from .serializers import TicketSerializer
consumer = Consumer(settings.TICKET_CREATION_TOPIC)
consumer.register(TicketSerializer)
return consumer
def test_missing_message_ignored(self, KafkaConsumer):
self.mock_consumer(KafkaConsumer, value=b'json:{"version":1}')
FakeStateSerializer = self.mock_state_serializer()
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.run(iter_limit=1)
self.assertEqual(FakeStateSerializer.call_count, 0)
def handle(self, *args, **options):
while True:
consumer = Consumer('test', consumer_timeout_ms=1000)
consumer.register(PrometheusSerializer)
consumer.run()
time.sleep(1)