def produce_test_messages_with_avro(
avro_producer: AvroProducer,
topic: Tuple[str, int]) -> Iterable[KafkaMessage]:
topic_name, num_partitions = topic
with open("tests/test_samples/key_schema.avsc", "r") as file:
key_schema = load_schema(file.read())
with open("tests/test_samples/value_schema.avsc", "r") as file:
value_schema = load_schema(file.read())
messages = []
for i in range(10):
partition = random.randrange(0, num_partitions)
key = {"id": str(i)}
value = {"first": "Firstname", "last": "Lastname"}
messages.append(
KafkaMessage(json.dumps(key), json.dumps(value), partition,
key_schema, value_schema))
avro_producer.produce(
topic=topic_name,
key=key,
value=value,
key_schema=key_schema,
value_schema=value_schema,
partition=partition,
)
avro_producer.flush()
return messages
def read_message_from_file(self) -> Iterable[KafkaMessage]:
while True:
try:
record = pickle.load(self.file)
except EOFError:
return
schema_directory = self.directory / record["schema_directory_name"]
key_schema = load_schema(
(schema_directory /
"key_schema.avsc").read_text(encoding="utf-8"))
value_schema = load_schema(
(schema_directory /
"value_schema.avsc").read_text(encoding="utf-8"))
yield KafkaMessage(json.dumps(record["key"]),
json.dumps(record["value"]),
record["partition"], key_schema, value_schema)
def mk_avro_schema(field_name: str, field_type: str) -> avro.schema.Schema:
return load_schema(
json.dumps({
"type": "record",
"namespace": "com.example",
"name": f"MySchema_{field_name}",
"fields": [{
"name": field_name,
"type": field_type
}],
}))