fix ut (#4455)

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

plugin/src/main/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerAction.java

@@ -117,7 +117,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
                 MemoryConfiguration configuration = container.getConfiguration();
                 String indexPrefix = configuration.getIndexPrefix();
 
-                memoryContainerHelper.countContainersWithPrefix(indexPrefix, tenantId, ActionListener.wrap(count -> {
+                memoryContainerHelper.countContainersWithPrefix(indexPrefix, ActionListener.wrap(count -> {
                     if (count > 1) {
                         // Multiple containers share this prefix - cannot delete indices
                         String error = String

plugin/src/main/java/org/opensearch/ml/action/memorycontainer/memory/TransportSearchMemoriesAction.java

@@ -118,9 +118,10 @@ private void searchMemories(
             memoryContainerHelper.addContainerIdFilter(input.getMemoryContainerId(), input.getSearchSourceBuilder());
 
             // Add owner filter for non-admin users
-            // if (!memoryContainerHelper.isAdminUser(user)) {
-            // memoryContainerHelper.addOwnerIdFilter(user, input.getSearchSourceBuilder());
-            // }
+            // TODO: now we don't support owner id in remote store, check if it's necessary to add
+            if (container.getConfiguration().getRemoteStore() == null && !memoryContainerHelper.isAdminUser(user)) {
+                memoryContainerHelper.addOwnerIdFilter(user, input.getSearchSourceBuilder());
+            }
 
             SearchRequest searchRequest = new SearchRequest(indexName).source(input.getSearchSourceBuilder());
             // TODO: add search pipeline support in SearchRequest

plugin/src/main/java/org/opensearch/ml/helper/MemoryContainerHelper.java

@@ -722,10 +722,11 @@ public SearchSourceBuilder addUserBackendRolesFilter(User user, SearchSourceBuil
     }
 
     public SearchSourceBuilder addOwnerIdFilter(User user, SearchSourceBuilder searchSourceBuilder) {
-        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
-        if (user != null) {
-            boolQueryBuilder.should(QueryBuilders.termsQuery(OWNER_ID_FIELD, user.getName()));
+        if (user == null) {
+            return searchSourceBuilder;
         }
+        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
+        boolQueryBuilder.should(QueryBuilders.termsQuery(OWNER_ID_FIELD, user.getName()));
         return applyFilterToSearchSource(searchSourceBuilder, boolQueryBuilder);
     }
 
@@ -799,10 +800,9 @@ public String getOwnerId(User user) {
      * Count memory containers with the specified index prefix
      *
      * @param indexPrefix the index prefix to search for
-     * @param tenantId the tenant ID (optional)
      * @param listener action listener returning the count
      */
-    public void countContainersWithPrefix(String indexPrefix, String tenantId, ActionListener<Long> listener) {
+    public void countContainersWithPrefix(String indexPrefix, ActionListener<Long> listener) {
         if (indexPrefix == null || indexPrefix.isBlank()) {
             listener.onResponse(0L);
             return;

plugin/src/test/java/org/opensearch/ml/action/memorycontainer/TransportDeleteMemoryContainerActionTests.java

@@ -393,10 +393,10 @@ public void testHandleDeleteResponse_WithDeleteAllMemories_Success() {
 
         // Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
         doAnswer(invocation -> {
-            ActionListener<Long> listener = invocation.getArgument(2);
+            ActionListener<Long> listener = invocation.getArgument(1);
             listener.onResponse(1L);
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Mock successful deleteIndex operation
         doAnswer(invocation -> {
@@ -466,10 +466,10 @@ public void testHandleDeleteResponse_WithDeleteAllMemories_DeleteIndexFailure()
 
         // Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
         doAnswer(invocation -> {
-            ActionListener<Long> listener = invocation.getArgument(2);
+            ActionListener<Long> listener = invocation.getArgument(1);
             listener.onResponse(1L);
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Mock failed deleteIndex operation
         doAnswer(invocation -> {
@@ -613,10 +613,10 @@ public void testDeleteMemoryContainer_WithSelectiveMemories_Success() {
 
         // Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
         doAnswer(invocation -> {
-            ActionListener<Long> listener = invocation.getArgument(2);
+            ActionListener<Long> listener = invocation.getArgument(1);
             listener.onResponse(1L);
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Mock successful deleteIndex operation
         doAnswer(invocation -> {
@@ -684,10 +684,10 @@ public void testDeleteMemoryContainer_WithSelectiveMemories_UnknownType() {
 
         // Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
         doAnswer(invocation -> {
-            ActionListener<Long> listener = invocation.getArgument(2);
+            ActionListener<Long> listener = invocation.getArgument(1);
             listener.onResponse(1L);
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Mock successful deleteIndex operation
         doAnswer(invocation -> {
@@ -789,10 +789,10 @@ public void testDeleteMemoryContainer_WithSharedPrefix_PreventsDeletion() {
 
         // Mock countContainersWithPrefix to return 2 (shared prefix - unsafe to delete indices)
         doAnswer(invocation -> {
-            ActionListener<Long> listener = invocation.getArgument(2);
+            ActionListener<Long> listener = invocation.getArgument(1);
             listener.onResponse(2L);
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Execute
         transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
@@ -847,7 +847,7 @@ public void testDeleteMemoryContainer_CountCheckFailure_AbortsOperation() {
             ActionListener<Long> listener = invocation.getArgument(2);
             listener.onFailure(new RuntimeException("Failed to count containers"));
             return null;
-        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
+        }).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
 
         // Execute
         transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
@@ -903,7 +903,7 @@ public void testDeleteMemoryContainer_WithoutIndexDeletion_SkipsCountCheck() {
         transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
 
         // Verify countContainersWithPrefix was NEVER called (no index deletion requested)
-        verify(memoryContainerHelper, timeout(1000).times(0)).countContainersWithPrefix(any(), any(), any());
+        verify(memoryContainerHelper, timeout(1000).times(0)).countContainersWithPrefix(any(), any());
 
         // Verify deleteIndex was NEVER called
         verify(memoryContainerHelper, timeout(1000).times(0)).deleteIndex(any(), any(), any());