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 convo1 = new StepListener<>(); + index.createConversation("Conversation1", "app", Map.of("k1", "v1"), convo1); + + StepListener convo2 = new StepListener<>(); + convo1.whenComplete(cid -> { index.createConversation("Mehul Conversation", convo2); }, e -> { + cdl.countDown(); + log.error(e); + assert (false); + }); + + StepListener 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;