add test for search conversation

Signed-off-by: Hailong Cui <[email protected]>
opensearch-project · Jul 26, 2024 · a82db1e · a82db1e
1 parent 8c2d61e
commit a82db1e
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/memory/src/test/java/org/opensearch/ml/memory/index/ConversationMetaIndexITTests.java b/memory/src/test/java/org/opensearch/ml/memory/index/ConversationMetaIndexITTests.java
@@ -21,6 +21,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.Stack;
 import java.util.concurrent.CountDownLatch;
@@ -671,4 +672,47 @@ public void testCanCreateConversationWithAdditionalInfo() {
             log.error(e);
         }
     }
+
+    public void testCanQueryOverConversationsByAdditionalInfo() {
+        CountDownLatch cdl = new CountDownLatch(1);
+        StepListener<String> convo1 = new StepListener<>();
+        index.createConversation("Conversation1", "app", Map.of("k1", "v1"), convo1);
+
+        StepListener<String> convo2 = new StepListener<>();
+        convo1.whenComplete(cid -> { index.createConversation("Mehul Conversation", convo2); }, e -> {
+            cdl.countDown();
+            log.error(e);
+            assert (false);
+        });
+
+        StepListener<SearchResponse> search = new StepListener<>();
+        convo2.whenComplete(cid -> {
+            SearchRequest request = new SearchRequest();
+            request.source(new SearchSourceBuilder());
+            request.source().query(QueryBuilders.matchQuery(ConversationalIndexConstants.META_ADDITIONAL_INFO_FIELD + ".k1", "v1"));
+            index.searchConversations(request, search);
+        }, e -> {
+            cdl.countDown();
+            log.error(e);
+            assert (false);
+        });
+
+        search.whenComplete(response -> {
+            log.info("SEARCH RESPONSE");
+            log.info(response.toString());
+            cdl.countDown();
+            assert (response.getHits().getAt(0).getId().equals(convo1.result()));
+            Assert.assertEquals(1L, Objects.requireNonNull(response.getHits().getTotalHits()).value);
+        }, e -> {
+            cdl.countDown();
+            log.error(e);
+            assert (false);
+        });
+
+        try {
+            cdl.await();
+        } catch (InterruptedException e) {
+            log.error(e);
+        }
+    }
 }
diff --git a/.../test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java b/.../test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java
@@ -31,6 +31,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Before;
 import org.mockito.ArgumentCaptor;