From 966495a2a97921f36ca845a108a1c185c2375730 Mon Sep 17 00:00:00 2001
From: Mike Siegel <msiegel@wayfair.com>
Date: Wed, 10 Oct 2018 10:04:30 -0400
Subject: [PATCH] PEP8 changes

---
 parsedmarc/kafkaclient.py | 42 +++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/parsedmarc/kafkaclient.py b/parsedmarc/kafkaclient.py
index b1e27730..7f48c649 100644
--- a/parsedmarc/kafkaclient.py
+++ b/parsedmarc/kafkaclient.py
@@ -2,18 +2,26 @@
 # -*- coding: utf-8 -*-
 
 from kafka import KafkaProducer
-from collections import OrderedDict
+from kafka.errors import NoBrokersAvailable, UnknownTopicOrPartitionError
 import json
 
+
 class KafkaError(RuntimeError):
     """Raised when a Kafka error occurs"""
 
-class KafkaClient(object):
-  def __init__(self, kafka_hosts):
-    self.producer = KafkaProducer(value_serializer=lambda v: json.dumps(v).encode('utf-8'),
-                                    bootstrap_servers=kafka_hosts)
 
-  def save_aggregate_reports_to_kafka(self, aggregate_reports, aggregate_topic):
+class KafkaClient(object):
+    def __init__(self, kafka_hosts):
+        try:
+            def serializer(v): lambda v: json.dumps(v).encode('utf-8')
+            self.producer = KafkaProducer(
+                          value_serializer=serializer,
+                          bootstrap_servers=kafka_hosts)
+        except NoBrokersAvailable:
+            raise KafkaError("No Kafka brokers availabe")
+
+    def save_aggregate_reports_to_kafka(self, aggregate_reports,
+                                        aggregate_topic):
         """
         Saves aggregate DMARC reports to Kafka
 
@@ -28,20 +36,18 @@ class KafkaClient(object):
         if len(aggregate_reports) < 1:
             return
 
-        for record in aggregate_reports['records']:
-            buffer = OrderedDict([('xml_schema', aggregate_reports['xml_schema']),
-                                ('report_metadata',aggregate_reports['report_metadata']),
-                                ('records', record)])
-            self.producer.send(aggregate_topic, buffer)
+        try:
+            self.producer.send(aggregate_topic, aggregate_reports)
+        except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
         self.producer.flush()
 
-
-  def save_forensic_reports_to_kafka(self, forensic_reports, forensic_topic):
+    def save_forensic_reports_to_kafka(self, forensic_reports, forensic_topic):
             """
             Saves forensic DMARC reports to Kafka
 
             Args:
-                forensic_reports (list):  A list of forensic report dictionaries
+                forensic_reports (list):  A list of forensic report dicts
                 to save to kafka
 
             """
@@ -51,6 +57,8 @@ class KafkaClient(object):
             if len(forensic_reports) < 1:
                 return
 
-            for report in forensic_reports:
-                self.producer.send(forensic_topic, json.dumps(report))
-                self.producer.flush()
+            try:
+                self.producer.send(forensic_topic, forensic_reports)
+            except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
+            self.producer.flush()