Skip to content

Commit e0df482

Browse files
ylwu-amznylwu-amzn
committed
fix more ut
Signed-off-by: Yaliang Wu <[email protected]>
1 parent 187961b commit e0df482

File tree

4 files changed

+51
-36
lines changed

4 files changed

+51
-36
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<Delete
117117
MemoryConfiguration configuration = container.getConfiguration();
118118
String indexPrefix = configuration.getIndexPrefix();
119119

120-
memoryContainerHelper.countContainersWithPrefix(indexPrefix, tenantId, ActionListener.wrap(count -> {
120+
memoryContainerHelper.countContainersWithPrefix(indexPrefix, ActionListener.wrap(count -> {
121121
if (count > 1) {
122122
// Multiple containers share this prefix - cannot delete indices
123123
String error = String

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -799,10 +799,9 @@ public String getOwnerId(User user) {
799799
* Count memory containers with the specified index prefix
800800
*
801801
* @param indexPrefix the index prefix to search for
802-
* @param tenantId the tenant ID (optional)
803802
* @param listener action listener returning the count
804803
*/
805-
public void countContainersWithPrefix(String indexPrefix, String tenantId, ActionListener<Long> listener) {
804+
public void countContainersWithPrefix(String indexPrefix, ActionListener<Long> listener) {
806805
if (indexPrefix == null || indexPrefix.isBlank()) {
807806
listener.onResponse(0L);
808807
return;

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -393,10 +393,10 @@ public void testHandleDeleteResponse_WithDeleteAllMemories_Success() {
393393

394394
// Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
395395
doAnswer(invocation -> {
396-
ActionListener<Long> listener = invocation.getArgument(2);
396+
ActionListener<Long> listener = invocation.getArgument(1);
397397
listener.onResponse(1L);
398398
return null;
399-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
399+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
400400

401401
// Mock successful deleteIndex operation
402402
doAnswer(invocation -> {
@@ -466,10 +466,10 @@ public void testHandleDeleteResponse_WithDeleteAllMemories_DeleteIndexFailure()
466466

467467
// Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
468468
doAnswer(invocation -> {
469-
ActionListener<Long> listener = invocation.getArgument(2);
469+
ActionListener<Long> listener = invocation.getArgument(1);
470470
listener.onResponse(1L);
471471
return null;
472-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
472+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
473473

474474
// Mock failed deleteIndex operation
475475
doAnswer(invocation -> {
@@ -613,10 +613,10 @@ public void testDeleteMemoryContainer_WithSelectiveMemories_Success() {
613613

614614
// Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
615615
doAnswer(invocation -> {
616-
ActionListener<Long> listener = invocation.getArgument(2);
616+
ActionListener<Long> listener = invocation.getArgument(1);
617617
listener.onResponse(1L);
618618
return null;
619-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
619+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
620620

621621
// Mock successful deleteIndex operation
622622
doAnswer(invocation -> {
@@ -684,10 +684,10 @@ public void testDeleteMemoryContainer_WithSelectiveMemories_UnknownType() {
684684

685685
// Mock countContainersWithPrefix to return 1 (sole owner - safe to delete indices)
686686
doAnswer(invocation -> {
687-
ActionListener<Long> listener = invocation.getArgument(2);
687+
ActionListener<Long> listener = invocation.getArgument(1);
688688
listener.onResponse(1L);
689689
return null;
690-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
690+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
691691

692692
// Mock successful deleteIndex operation
693693
doAnswer(invocation -> {
@@ -789,10 +789,10 @@ public void testDeleteMemoryContainer_WithSharedPrefix_PreventsDeletion() {
789789

790790
// Mock countContainersWithPrefix to return 2 (shared prefix - unsafe to delete indices)
791791
doAnswer(invocation -> {
792-
ActionListener<Long> listener = invocation.getArgument(2);
792+
ActionListener<Long> listener = invocation.getArgument(1);
793793
listener.onResponse(2L);
794794
return null;
795-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
795+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
796796

797797
// Execute
798798
transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
@@ -847,7 +847,7 @@ public void testDeleteMemoryContainer_CountCheckFailure_AbortsOperation() {
847847
ActionListener<Long> listener = invocation.getArgument(2);
848848
listener.onFailure(new RuntimeException("Failed to count containers"));
849849
return null;
850-
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any(), any());
850+
}).when(memoryContainerHelper).countContainersWithPrefix(any(), any());
851851

852852
// Execute
853853
transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
@@ -903,7 +903,7 @@ public void testDeleteMemoryContainer_WithoutIndexDeletion_SkipsCountCheck() {
903903
transportDeleteMemoryContainerAction.doExecute(null, deleteRequest, actionListener);
904904

905905
// Verify countContainersWithPrefix was NEVER called (no index deletion requested)
906-
verify(memoryContainerHelper, timeout(1000).times(0)).countContainersWithPrefix(any(), any(), any());
906+
verify(memoryContainerHelper, timeout(1000).times(0)).countContainersWithPrefix(any(), any());
907907

908908
// Verify deleteIndex was NEVER called
909909
verify(memoryContainerHelper, timeout(1000).times(0)).deleteIndex(any(), any(), any());

‎plugin/src/test/java/org/opensearch/ml/helper/MemoryContainerHelperTests.java‎

Lines changed: 37 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,6 @@
6161
import org.opensearch.ml.common.memorycontainer.MemoryStrategy;
6262
import org.opensearch.ml.common.memorycontainer.MemoryType;
6363
import org.opensearch.remote.metadata.client.GetDataObjectResponse;
64-
import org.opensearch.remote.metadata.client.SearchDataObjectResponse;
6564
import org.opensearch.search.SearchHit;
6665
import org.opensearch.search.SearchHits;
6766
import org.opensearch.search.aggregations.InternalAggregations;
@@ -253,16 +252,21 @@ public void testSearchData() {
253252
request.source(new SearchSourceBuilder());
254253

255254
SearchResponse searchResponse = createSearchResponse(2);
256-
SearchDataObjectResponse response = new SearchDataObjectResponse(searchResponse);
257-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(CompletableFuture.completedFuture(response));
255+
doAnswer(invocation -> {
256+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
257+
listener.onResponse(searchResponse);
258+
return null;
259+
}).when(client).search(any(), any());
258260

259261
PlainActionFuture<SearchResponse> future = PlainActionFuture.newFuture();
260262
helper.searchData(configuration, request, future);
261263
assertSame(searchResponse, future.actionGet());
262264

263-
CompletableFuture<SearchDataObjectResponse> failed = new CompletableFuture<>();
264-
failed.completeExceptionally(new RuntimeException("search failure"));
265-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(failed);
265+
doAnswer(invocation -> {
266+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
267+
listener.onFailure(new RuntimeException("search failure"));
268+
return null;
269+
}).when(client).search(any(), any());
266270

267271
PlainActionFuture<SearchResponse> failure = PlainActionFuture.newFuture();
268272
helper.searchData(configuration, request, failure);
@@ -277,17 +281,23 @@ public void testSearchDataWithSystemIndex() {
277281
request.source(new SearchSourceBuilder());
278282

279283
SearchResponse searchResponse = createSearchResponse(3);
280-
SearchDataObjectResponse response = new SearchDataObjectResponse(searchResponse);
281-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(CompletableFuture.completedFuture(response));
284+
285+
doAnswer(invocation -> {
286+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
287+
listener.onResponse(searchResponse);
288+
return null;
289+
}).when(client).search(any(), any());
282290

283291
PlainActionFuture<SearchResponse> future = PlainActionFuture.newFuture();
284292
helper.searchData(systemConfig, request, future);
285293
assertSame(searchResponse, future.actionGet());
286294

287295
// Test failure with system index
288-
CompletableFuture<SearchDataObjectResponse> failed = new CompletableFuture<>();
289-
failed.completeExceptionally(new RuntimeException("system index search failure"));
290-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(failed);
296+
doAnswer(invocation -> {
297+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
298+
listener.onFailure(new RuntimeException("system index search failure"));
299+
return null;
300+
}).when(client).search(any(), any());
291301

292302
PlainActionFuture<SearchResponse> failure = PlainActionFuture.newFuture();
293303
helper.searchData(systemConfig, request, failure);
@@ -642,33 +652,39 @@ public void testGetLlmResultPath() {
642652
public void testCountContainersWithPrefix() {
643653
// Test with null prefix
644654
PlainActionFuture<Long> nullFuture = PlainActionFuture.newFuture();
645-
helper.countContainersWithPrefix(null, null, nullFuture);
655+
helper.countContainersWithPrefix(null, nullFuture);
646656
assertEquals(0L, nullFuture.actionGet().longValue());
647657

648658
// Test with blank prefix
649659
PlainActionFuture<Long> blankFuture = PlainActionFuture.newFuture();
650-
helper.countContainersWithPrefix("", null, blankFuture);
660+
helper.countContainersWithPrefix("", blankFuture);
651661
assertEquals(0L, blankFuture.actionGet().longValue());
652662

653663
SearchResponse searchResponse = createSearchResponse(3);
654-
SearchDataObjectResponse response = new SearchDataObjectResponse(searchResponse);
655-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(CompletableFuture.completedFuture(response));
664+
665+
doAnswer(invocation -> {
666+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
667+
listener.onResponse(searchResponse);
668+
return null;
669+
}).when(client).search(any(), any());
656670

657671
PlainActionFuture<Long> future = PlainActionFuture.newFuture();
658-
helper.countContainersWithPrefix("prefix", null, future);
672+
helper.countContainersWithPrefix("prefix", future);
659673
assertEquals(3L, future.actionGet().longValue());
660674

661675
// Test with tenant ID
662676
PlainActionFuture<Long> tenantFuture = PlainActionFuture.newFuture();
663-
helper.countContainersWithPrefix("prefix", "tenant123", tenantFuture);
677+
helper.countContainersWithPrefix("prefix", tenantFuture);
664678
assertEquals(3L, tenantFuture.actionGet().longValue());
665679

666-
CompletableFuture<SearchDataObjectResponse> failed = new CompletableFuture<>();
667-
failed.completeExceptionally(new RuntimeException("fail"));
668-
when(sdkClient.searchDataObjectAsync(any())).thenReturn(failed);
680+
doAnswer(invocation -> {
681+
ActionListener<SearchResponse> listener = (ActionListener<SearchResponse>) invocation.getArguments()[1];
682+
listener.onFailure(new RuntimeException("fail"));
683+
return null;
684+
}).when(client).search(any(), any());
669685

670686
PlainActionFuture<Long> failure = PlainActionFuture.newFuture();
671-
helper.countContainersWithPrefix("prefix", null, failure);
687+
helper.countContainersWithPrefix("prefix", failure);
672688
RuntimeException exception = expectThrows(RuntimeException.class, failure::actionGet);
673689
assertEquals("fail", exception.getMessage());
674690
}

0 commit comments

Comments
 (0)