1 files changed, 21 insertions, 16 deletions

diff --git a/netconfsimulator/src/main/java/org/onap/netconfsimulator/kafka/StoreService.java b/netconfsimulator/src/main/java/org/onap/netconfsimulator/kafka/StoreService.java
index 5fddff5..6bd8390 100644
--- a/netconfsimulator/src/main/java/org/onap/netconfsimulator/kafka/StoreService.java
+++ b/netconfsimulator/src/main/java/org/onap/netconfsimulator/kafka/StoreService.java
@@ -7,9 +7,9 @@
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  *      http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -28,6 +28,7 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.kafka.core.ConsumerFactory;
 import org.springframework.stereotype.Service;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -48,33 +49,33 @@ public class StoreService {
         this.consumerFactory = consumerFactory;
     }
 
-    List<MessageDTO> getAllMessages() {
-        List<MessageDTO> messages = new ArrayList<>();
-        String clientID = Long.toString(Instant.now().getEpochSecond());
-        try (Consumer<String, String> consumer = consumerFactory.createConsumer(clientID, clientID)) {
+    List<Message> getAllMessages() {
+        List<Message> messages = new ArrayList<>();
+        String clientId = Long.toString(Instant.now().getEpochSecond());
+        try (Consumer<String, String> consumer = consumerFactory.createConsumer(clientId, clientId)) {
             consumer.subscribe(TOPICS_TO_SUBSCRIBE);
-            ConsumerRecords<String, String> consumerRecords = consumer.poll(CONSUMING_DURATION_IN_MS);
+            ConsumerRecords<String, String> consumerRecords = pollConsumerRecords(consumer);
             consumerRecords.forEach(
                 consumerRecord ->
-                    messages.add(new MessageDTO(consumerRecord.timestamp(), consumerRecord.value())));
+                    messages.add(new Message(consumerRecord.timestamp(), consumerRecord.value())));
             log.debug(String.format("consumed %d messages", consumerRecords.count()));
-            }
+        }
         return messages;
     }
 
-    List<MessageDTO> getLastMessages(long offset) {
-        List<MessageDTO> messages = new ArrayList<>();
+    List<Message> getLastMessages(long offset) {
+        List<Message> messages = new ArrayList<>();
         try (Consumer<String, String> consumer = createConsumer(offset)) {
-            ConsumerRecords<String, String> consumerRecords = consumer.poll(CONSUMING_DURATION_IN_MS);
+            ConsumerRecords<String, String> consumerRecords = pollConsumerRecords(consumer);
             consumerRecords.forEach(consumerRecord ->
-                    messages.add(new MessageDTO(consumerRecord.timestamp(), consumerRecord.value())));
+                messages.add(new Message(consumerRecord.timestamp(), consumerRecord.value())));
         }
         return messages;
     }
 
     private Consumer<String, String> createConsumer(long offsetFromLastIndex) {
-        String clientID = Long.toString(Instant.now().getEpochSecond());
-        Consumer<String, String> consumer = consumerFactory.createConsumer(clientID, clientID);
+        String clientId = Long.toString(Instant.now().getEpochSecond());
+        Consumer<String, String> consumer = consumerFactory.createConsumer(clientId, clientId);
         consumer.subscribe(TOPICS_TO_SUBSCRIBE);
         seekConsumerTo(consumer, offsetFromLastIndex);
         return consumer;
@@ -82,10 +83,14 @@ public class StoreService {
 
     private void seekConsumerTo(Consumer<String, String> consumer, long offsetFromLastIndex) {
         consumer.seekToEnd(consumer.assignment());
-        consumer.poll(CONSUMING_DURATION_IN_MS);
+        pollConsumerRecords(consumer);
         TopicPartition topicPartition = consumer.assignment().iterator().next();
         long topicCurrentSize = consumer.position(topicPartition);
         long indexToSeek = offsetFromLastIndex > topicCurrentSize ? 0 : topicCurrentSize - offsetFromLastIndex;
         consumer.seek(topicPartition, indexToSeek);
     }
+
+    private ConsumerRecords<String, String> pollConsumerRecords(Consumer<String, String> consumer) {
+        return consumer.poll(Duration.ofMillis(CONSUMING_DURATION_IN_MS));
+    }
 }