def test_kafka_key_deserializer_serializer(mock_kafka_consumer, mock_kafka_producer):
create_kafka_consumer(BOOTSTRAP_SERVERS, GROUP_ID, [TOPIC])
create_kafka_producer(BOOTSTRAP_SERVERS)
key_deserializer = get_key_deserializer(mock_kafka_consumer)
key_serializer = get_key_serializer(mock_kafka_producer)
assert CLIENT_ID == key_deserializer(key_serializer(CLIENT_ID))
def test_kafka_producer_pass_through(mock_kafka_producer):
create_kafka_producer(BOOTSTRAP_SERVERS)
mock_kafka_producer.assert_called_with(
bootstrap_servers=BOOTSTRAP_SERVERS,
value_serializer=ANY,
key_serializer=ANY,
connections_max_idle_ms=MAX_IDLE_TIME_MS,
)
def test_kafka_client_create_consumer_producer(mocker, mock_kafka_consumer, mock_kafka_producer, kafka_config):
mock_create_topics = mocker.patch.object(bai_kafka_utils.kafka_client, "create_kafka_topics")
consumer, producer = create_kafka_consumer_producer(kafka_config, SERVICE_NAME)
expected_consumer = create_kafka_consumer(
kafka_config.bootstrap_servers,
kafka_config.consumer_group_id,
[kafka_config.consumer_topic, kafka_config.cmd_submit_topic],
)
expected_producer = create_kafka_producer(kafka_config.bootstrap_servers)
mock_create_topics.assert_called_once()
assert consumer == expected_consumer
assert producer == expected_producer
def main(argv=None):
configure_logging(level=logging.INFO)
logger = logging.getLogger("BAI_CRON_JOB")
logger.info("Loading configuration")
logger.info(os.environ)
try:
config = get_config(argv, os.environ)
logging.info("Updating benchmark event's message and action id")
scheduled_benchmark_event = dacite.from_dict(
data_class=FetcherBenchmarkEvent, data=config.benchmark_event)
benchmark_event = create_benchmark_event(scheduled_benchmark_event)
logging.info("Creating Kafka producer")
kafka_producer = create_kafka_producer(config.kafka_bootstrap_servers)
event_emitter = EventEmitter(
name=SERVICE_NAME,
version=__version__,
pod_name=get_pod_name(),
status_topic=config.status_topic,
kakfa_producer=kafka_producer,
)
logging.info(
f"Submitting benchmark with action id {benchmark_event.action_id}")
event_emitter.send_event(event=benchmark_event,
topic=config.producer_topic)
# Send the message against the original scheduled benchmark event so it is logged
# against the original action_id
event_emitter.send_status_message_event(
handled_event=scheduled_benchmark_event,
status=Status.SUCCEEDED,
msg=
f"Spawning benchmark with action_id: [{benchmark_event.action_id}]",
)
logging.debug("Closing producer")
kafka_producer.close()
except Exception as err:
logger.exception(f"Fatal error submitting benchmark job: {err}")
sys.exit(1)
logging.info("Success")
def __init__(self,
action_id: str,
client_id: str,
labels: Dict[str, str],
*,
topic: str,
bootstrap_servers: List[str] = None,
key: str = None):
self.labels = {
"action-id": action_id,
"client-id": client_id,
**labels
}
if bootstrap_servers is None:
bootstrap_servers = ["localhost:9092"]
self._producer = create_kafka_producer(bootstrap_servers)
self._key = key
self._topic = topic
def kafka_producer_to_consume(kafka_service_config: KafkaServiceConfig):
print("Creating a producer...\n")
kafka_producer = create_kafka_producer(kafka_service_config.bootstrap_servers)
yield kafka_producer
print("Closing producer...\n")
kafka_producer.close()
def test_kafka_key_serializer(mock_kafka_producer):
create_kafka_producer(BOOTSTRAP_SERVERS)
key_serializer = get_key_serializer(mock_kafka_producer)
assert CLIENT_ID_SERIALIZED == key_serializer(CLIENT_ID)