feat: allow user to set callback for simple produce call (#23)

* feat: allow user to set callback for simple produce call

aet/kafka_utils.py

@@ -83,8 +83,10 @@ def get_broker_info(kclient, scope='all'):
                 else:
                     errstr = ""
 
-                t_str.append("partition {} leader: {}, replicas: {}, isrs: {}".format(
-                    p.id, p.leader, p.replicas, p.isrs, errstr))
+                t_str.append(
+                    f"partition {p.id} leader: {p.leader}, "
+                    f"replicas: {p.replicas}, isrs: {p.isrs}, err: {errstr}"
+                )
             res['topics'].append(t_str)
         if scope in res.keys():
             return res[scope]
@@ -143,8 +145,9 @@ def kafka_callback(err=None, msg=None, _=None, **kwargs):
                 LOG.error(f'NO-SAVE: {_id} in | err {err.name()}')
 
 
-def produce(docs, schema, topic_name, producer):
-
+def produce(docs, schema, topic_name, producer, callback=None):
+    if not callback:
+        callback = kafka_callback
     with io.BytesIO() as bytes_writer:
         writer = DataFileWriter(
             bytes_writer, DatumWriter(), schema, codec='deflate')
@@ -159,15 +162,14 @@ def produce(docs, schema, topic_name, producer):
                 # Message doesn't have the proper format for the current schema.
                 LOG.debug(
                     f"SCHEMA_MISMATCH:NOT SAVED! TOPIC:{topic_name}, ID:{_id}")
-
         writer.flush()
         raw_bytes = bytes_writer.getvalue()
 
     producer.poll(0)
     producer.produce(
         topic_name,
         raw_bytes,
-        callback=kafka_callback,
+        callback=callback,
         headers={
             'avro_size': str(len(_ids)),
             'contains_id': json.dumps(_ids)

tests/__init__.py

@@ -615,6 +615,7 @@ def mocked_consumer():
     consumer.stop()
     sleep(.5)
 
+
 # API Assets
 @pytest.mark.unit
 @pytest.fixture(scope="module")