Add OpenSearhMemoryHandler unit tests for singular get

Signed-off-by: HenryL27 <[email protected]>
HenryL27 · Oct 11, 2023 · f148f52 · f148f52
1 parent b60c0c2
commit f148f52
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/memory/src/test/java/org/opensearch/ml/memory/index/InteractionsIndexTests.java b/memory/src/test/java/org/opensearch/ml/memory/index/InteractionsIndexTests.java
@@ -63,9 +63,6 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.SendRequestTransportException;
 
-import lombok.extern.log4j.Log4j2;
-
-@Log4j2
 public class InteractionsIndexTests extends OpenSearchTestCase {
     @Mock
     Client client;

diff --git a/.../test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java b/.../test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java
@@ -25,6 +25,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.time.Instant;
 import java.util.List;
 
 import org.junit.Before;
@@ -268,4 +269,26 @@ public void testSearchInteractions_Future() {
         ActionFuture<SearchResponse> result = cmHandler.searchInteractions(cid, request);
         assert (result.actionGet().equals(response));
     }
+
+    public void testGetAConversation_Future() {
+        ConversationMeta response = new ConversationMeta("cid", Instant.now(), "boring name", null);
+        doAnswer(invocation -> {
+            ActionListener<ConversationMeta> listener = invocation.getArgument(1);
+            listener.onResponse(response);
+            return null;
+        }).when(conversationMetaIndex).getConversation(any(), any());
+        ActionFuture<ConversationMeta> result = cmHandler.getConversation("cid");
+        assert (result.actionGet().equals(response));
+    }
+
+    public void testGetAnInteraction_Future() {
+        Interaction interaction = new Interaction("iid", Instant.now(), "cid", "inp", "pt", "rsp", "ogn", "extra");
+        doAnswer(invocation -> {
+            ActionListener<Interaction> listener = invocation.getArgument(2);
+            listener.onResponse(interaction);
+            return null;
+        }).when(interactionsIndex).getInteraction(any(), any(), any());
+        ActionFuture<Interaction> result = cmHandler.getInteraction("cid", "iid");
+        assert (result.actionGet().equals(interaction));
+    }
 }