1 files changed, 15 insertions, 14 deletions

diff --git a/netconfsimulator/src/test/java/org/onap/netconfsimulator/kafka/StoreControllerTest.java b/netconfsimulator/src/test/java/org/onap/netconfsimulator/kafka/StoreControllerTest.java
index 02eec12..bd6fd14 100644
--- a/netconfsimulator/src/test/java/org/onap/netconfsimulator/kafka/StoreControllerTest.java
+++ b/netconfsimulator/src/test/java/org/onap/netconfsimulator/kafka/StoreControllerTest.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.
@@ -22,6 +22,7 @@ package org.onap.netconfsimulator.kafka;
 
 import java.time.Instant;
 import java.util.List;
+
 import org.assertj.core.api.Assertions;
 import org.assertj.core.util.Lists;
 import org.junit.Test;
@@ -39,9 +40,9 @@ public class StoreControllerTest {
     private static final String MESSAGE_2 = "message 2";
     private static final String MESSAGE_1 = "message 1";
 
-    private static final List<MessageDTO> ALL_MESSAGES = Lists.newArrayList(new MessageDTO(Instant.now().getEpochSecond(), MESSAGE_1),
-            new MessageDTO(Instant.now().getEpochSecond(), MESSAGE_2),
-            new MessageDTO(Instant.now().getEpochSecond(), MESSAGE_3));
+    private static final List<Message> ALL_MESSAGES = Lists.newArrayList(new Message(Instant.now().getEpochSecond(), MESSAGE_1),
+        new Message(Instant.now().getEpochSecond(), MESSAGE_2),
+        new Message(Instant.now().getEpochSecond(), MESSAGE_3));
 
     @Mock
     private StoreService service;
@@ -54,33 +55,33 @@ public class StoreControllerTest {
     public void lessShouldTakeAllMessagesTest() {
         when(service.getLastMessages(3)).thenReturn(ALL_MESSAGES);
 
-        List<MessageDTO> lessResponse = storeController.less(3);
+        List<Message> lessResponse = storeController.less(3);
 
         assertResponseContainsExpectedMessages(lessResponse, 3, MESSAGE_1, MESSAGE_2, MESSAGE_3);
     }
 
     @Test
     public void lessShouldTakeTwoMessagesTest() {
-        when(service.getLastMessages(2)).thenReturn(Lists.newArrayList(new MessageDTO(Instant.now().getEpochSecond(), MESSAGE_1)));
+        when(service.getLastMessages(2)).thenReturn(Lists.newArrayList(new Message(Instant.now().getEpochSecond(), MESSAGE_1)));
 
-        List<MessageDTO> lessResult = storeController.less(2);
+        List<Message> lessResult = storeController.less(2);
 
         assertResponseContainsExpectedMessages(lessResult, 1, MESSAGE_1);
     }
 
     @Test
-    public void shouldGetAllMessages(){
+    public void shouldGetAllMessages() {
         when(service.getAllMessages()).thenReturn(ALL_MESSAGES);
 
-        List<MessageDTO> allMsgResult = storeController.getAllConfigurationChanges();
+        List<Message> allMsgResult = storeController.getAllConfigurationChanges();
 
         assertResponseContainsExpectedMessages(allMsgResult, 3, MESSAGE_1, MESSAGE_2, MESSAGE_3);
     }
 
-    private void assertResponseContainsExpectedMessages(List<MessageDTO> actualMessages, int expectedMessageCount, String... expectedMessages){
-        Assertions.assertThat(actualMessages.stream().map(MessageDTO::getConfiguration))
-                .hasSize(expectedMessageCount)
-                .containsExactly(expectedMessages);
+    private void assertResponseContainsExpectedMessages(List<Message> actualMessages, int expectedMessageCount, String... expectedMessages) {
+        Assertions.assertThat(actualMessages.stream().map(Message::getConfiguration))
+            .hasSize(expectedMessageCount)
+            .containsExactly(expectedMessages);
     }
 
 }