From a82db1e9e40b249cdb837b2095b23129e3f8ef76 Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Fri, 26 Jul 2024 14:52:55 +0800
Subject: [PATCH] add test for search conversation

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 .../index/ConversationMetaIndexITTests.java   | 44 +++++++++++++++++++
 ...earchConversationalMemoryHandlerTests.java |  1 +
 2 files changed, 45 insertions(+)

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
index 09bb03eb61..5baefa358d 100644
--- 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/memory/src/test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java b/memory/src/test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java
index 0fe96e90ce..903be08338 100644
--- a/memory/src/test/java/org/opensearch/ml/memory/index/OpenSearchConversationalMemoryHandlerTests.java
+++ b/memory/src/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;