ATLAS-1908: updated NotificationConsumer with a method removed earlier to avoid…

notification/src/main/java/org/apache/atlas/kafka/AtlasKafkaConsumer.java

@@ -52,9 +52,15 @@ public class AtlasKafkaConsumer<T> extends AbstractNotificationConsumer<T> {
     }
 
     public List<AtlasKafkaMessage<T>> receive() {
+        return this.receive(this.pollTimeoutMilliSeconds);
+    }
+
+    @Override
+    public List<AtlasKafkaMessage<T>> receive(long timeoutMilliSeconds) {
+
         List<AtlasKafkaMessage<T>> messages = new ArrayList();
 
-        ConsumerRecords<?, ?> records = kafkaConsumer.poll(pollTimeoutMilliSeconds);
+        ConsumerRecords<?, ?> records = kafkaConsumer.poll(timeoutMilliSeconds);
 
         if (records != null) {
             for (ConsumerRecord<?, ?> record : records) {
@@ -70,6 +76,7 @@ public class AtlasKafkaConsumer<T> extends AbstractNotificationConsumer<T> {
         }
 
         return messages;
+
     }
 
 

notification/src/main/java/org/apache/atlas/notification/NotificationConsumer.java

@@ -44,4 +44,13 @@ public interface NotificationConsumer<T> {
      * @return List containing kafka message and partionId and offset.
      */
     List<AtlasKafkaMessage<T>> receive();
+
+    /**
+     * Fetch data for the topics from Kafka
+     * @param timeoutMilliSeconds poll timeout
+     * @return List containing kafka message and partionId and offset.
+     */
+    List<AtlasKafkaMessage<T>> receive(long timeoutMilliSeconds);
+
+
 }

notification/src/test/java/org/apache/atlas/notification/AbstractNotificationConsumerTest.java

@@ -204,6 +204,11 @@ public class AbstractNotificationConsumerTest {
 
         @Override
         public List<AtlasKafkaMessage<T>> receive() {
+            return receive(1000L);
+        }
+
+        @Override
+        public List<AtlasKafkaMessage<T>> receive(long timeoutMilliSeconds) {
             List<AtlasKafkaMessage<T>> tempMessageList = new ArrayList();
             for(Object json :  messageList) {
                 tempMessageList.add(new AtlasKafkaMessage(deserializer.deserialize((String)json), -1, -1));