From 1fcdf527e9e2f32a15c6f2e6b5c79e42b6a5f24b Mon Sep 17 00:00:00 2001 From: Ashish Singh Date: Tue, 10 Sep 2024 15:53:16 +0530 Subject: [PATCH] Add UTs Signed-off-by: Ashish Singh --- .../repositories/s3/StatsMetricPublisher.java | 4 + .../s3/async/S3AsyncDeleteHelper.java | 18 +- .../s3/S3BlobStoreContainerTests.java | 94 ++++++- .../s3/async/S3AsyncDeleteHelperTests.java | 236 ++++++++++++++++++ 4 files changed, 338 insertions(+), 14 deletions(-) create mode 100644 plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelperTests.java diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/StatsMetricPublisher.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/StatsMetricPublisher.java index 8d2772d42ebca..9f73c67df3b18 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/StatsMetricPublisher.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/StatsMetricPublisher.java @@ -95,6 +95,10 @@ public void publish(MetricCollection metricCollection) { public void close() {} }; + public MetricPublisher getDeleteObjectsMetricPublisher() { + return deleteObjectsMetricPublisher; + } + public MetricPublisher getObjectMetricPublisher = new MetricPublisher() { @Override public void publish(MetricCollection metricCollection) { diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelper.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelper.java index 5ce2c242bdefb..eed95c0e68ef1 100644 --- a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelper.java +++ b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelper.java @@ -50,11 +50,7 @@ static List> createDeleteBatches(List keys, int bulkDeleteS return batches; } - private static CompletableFuture executeDeleteBatches( - S3AsyncClient s3AsyncClient, - S3BlobStore blobStore, - List> batches - ) { + static CompletableFuture executeDeleteBatches(S3AsyncClient s3AsyncClient, S3BlobStore blobStore, List> batches) { CompletableFuture allDeletesFuture = CompletableFuture.completedFuture(null); for (List batch : batches) { @@ -64,16 +60,12 @@ private static CompletableFuture executeDeleteBatches( return allDeletesFuture; } - private static CompletableFuture executeSingleDeleteBatch( - S3AsyncClient s3AsyncClient, - S3BlobStore blobStore, - List batch - ) { + static CompletableFuture executeSingleDeleteBatch(S3AsyncClient s3AsyncClient, S3BlobStore blobStore, List batch) { DeleteObjectsRequest deleteRequest = bulkDelete(blobStore.bucket(), batch, blobStore); return s3AsyncClient.deleteObjects(deleteRequest).thenApply(S3AsyncDeleteHelper::processDeleteResponse); } - private static Void processDeleteResponse(DeleteObjectsResponse deleteObjectsResponse) { + static Void processDeleteResponse(DeleteObjectsResponse deleteObjectsResponse) { if (!deleteObjectsResponse.errors().isEmpty()) { logger.warn( () -> new ParameterizedMessage( @@ -88,7 +80,7 @@ private static Void processDeleteResponse(DeleteObjectsResponse deleteObjectsRes return null; } - private static DeleteObjectsRequest bulkDelete(String bucket, List blobs, S3BlobStore blobStore) { + static DeleteObjectsRequest bulkDelete(String bucket, List blobs, S3BlobStore blobStore) { return DeleteObjectsRequest.builder() .bucket(bucket) .delete( @@ -97,7 +89,7 @@ private static DeleteObjectsRequest bulkDelete(String bucket, List blobs .quiet(true) .build() ) - .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().deleteObjectsMetricPublisher)) + .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().getDeleteObjectsMetricPublisher())) .build(); } } diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java index 654d8a72690c4..ffe99dcb5e2ab 100644 --- a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java @@ -69,6 +69,7 @@ import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.model.UploadPartResponse; import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable; +import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Publisher; import org.opensearch.action.LatchedActionListener; import org.opensearch.common.blobstore.BlobContainer; @@ -101,16 +102,19 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -1275,6 +1279,94 @@ public void testTransformResponseToInputStreamContainer() throws Exception { assertEquals(inputStream.available(), inputStreamContainer.getInputStream().available()); } + public void testDeleteAsync() throws Exception { + final String bucketName = randomAlphaOfLengthBetween(1, 10); + final BlobPath blobPath = new BlobPath(); + + final S3BlobStore blobStore = mock(S3BlobStore.class); + when(blobStore.bucket()).thenReturn(bucketName); + when(blobStore.getStatsMetricPublisher()).thenReturn(new StatsMetricPublisher()); + when(blobStore.getBulkDeletesSize()).thenReturn(1); + + final S3AsyncClient s3AsyncClient = mock(S3AsyncClient.class); + final AmazonAsyncS3Reference asyncClientReference = mock(AmazonAsyncS3Reference.class); + when(blobStore.asyncClientReference()).thenReturn(asyncClientReference); + AmazonAsyncS3WithCredentials amazonAsyncS3WithCredentials = AmazonAsyncS3WithCredentials.create( + s3AsyncClient, + s3AsyncClient, + s3AsyncClient, + null + ); + when(asyncClientReference.get()).thenReturn(amazonAsyncS3WithCredentials); + + final List s3Objects = Arrays.asList( + S3Object.builder().key("key1").size(100L).build(), + S3Object.builder().key("key2").size(200L).build(), + S3Object.builder().key("key3").size(300L).build() + ); + + final ListObjectsV2Response response1 = ListObjectsV2Response.builder().contents(s3Objects.subList(0, 2)).build(); + final ListObjectsV2Response response2 = ListObjectsV2Response.builder().contents(s3Objects.subList(2, 3)).build(); + + final ListObjectsV2Publisher listPublisher = mock(ListObjectsV2Publisher.class); + AtomicInteger counter = new AtomicInteger(); + doAnswer(invocation -> { + Subscriber subscriber = invocation.getArgument(0); + subscriber.onSubscribe(new Subscription() { + @Override + public void request(long n) { + int currentCounter = counter.getAndIncrement(); + if (currentCounter == 0) { + subscriber.onNext(response1); + } + if (currentCounter == 1) { + subscriber.onNext(response2); + } + if (currentCounter == 2) { + subscriber.onComplete(); + } + } + + @Override + public void cancel() {} + }); + return null; + }).when(listPublisher).subscribe(any(Subscriber.class)); + + when(s3AsyncClient.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listPublisher); + + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn( + CompletableFuture.completedFuture(DeleteObjectsResponse.builder().build()) + ); + + final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference deleteResultRef = new AtomicReference<>(); + blobContainer.deleteAsync(new ActionListener<>() { + @Override + public void onResponse(DeleteResult deleteResult) { + deleteResultRef.set(deleteResult); + latch.countDown(); + } + + @Override + public void onFailure(Exception e) { + logger.error("exception during deleteAsync", e); + fail("Unexpected failure: " + e.getMessage()); + } + }); + + latch.await(); + + DeleteResult deleteResult = deleteResultRef.get(); + assertEquals(3, deleteResult.blobsDeleted()); + assertEquals(600, deleteResult.bytesDeleted()); + + verify(s3AsyncClient, times(1)).listObjectsV2Paginator(any(ListObjectsV2Request.class)); + verify(s3AsyncClient, times(3)).deleteObjects(any(DeleteObjectsRequest.class)); + } + private void mockObjectResponse(S3AsyncClient s3AsyncClient, String bucketName, String blobName, int objectSize) { final InputStream inputStream = new ByteArrayInputStream(randomByteArrayOfLength(objectSize)); diff --git a/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelperTests.java b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelperTests.java new file mode 100644 index 0000000000000..d7f924e05cc70 --- /dev/null +++ b/plugins/repository-s3/src/test/java/org/opensearch/repositories/s3/async/S3AsyncDeleteHelperTests.java @@ -0,0 +1,236 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.repositories.s3.async; + +import software.amazon.awssdk.metrics.MetricPublisher; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest; +import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse; +import software.amazon.awssdk.services.s3.model.ObjectIdentifier; +import software.amazon.awssdk.services.s3.model.S3Error; + +import org.opensearch.repositories.s3.S3BlobStore; +import org.opensearch.repositories.s3.StatsMetricPublisher; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.stream.Collectors; + +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class S3AsyncDeleteHelperTests extends OpenSearchTestCase { + + @Mock + private S3AsyncClient s3AsyncClient; + + @Mock + private S3BlobStore blobStore; + + @Override + public void setUp() throws Exception { + super.setUp(); + MockitoAnnotations.openMocks(this); + } + + public void testExecuteDeleteChain() { + List objectsToDelete = Arrays.asList("key1", "key2", "key3"); + CompletableFuture currentChain = CompletableFuture.completedFuture(null); + + // Mock the deleteObjects method of S3AsyncClient + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn( + CompletableFuture.completedFuture(DeleteObjectsResponse.builder().build()) + ); + + // Mock the getBulkDeletesSize method of S3BlobStore + when(blobStore.getBulkDeletesSize()).thenReturn(2); + + // Mock the getStatsMetricPublisher method of S3BlobStore to return a non-null value + StatsMetricPublisher mockMetricPublisher = mock(StatsMetricPublisher.class); + MetricPublisher mockDeleteObjectsMetricPublisher = mock(MetricPublisher.class); + when(blobStore.getStatsMetricPublisher()).thenReturn(mockMetricPublisher); + when(mockMetricPublisher.getDeleteObjectsMetricPublisher()).thenReturn(mockDeleteObjectsMetricPublisher); + + CompletableFuture newChain = S3AsyncDeleteHelper.executeDeleteChain( + s3AsyncClient, + blobStore, + objectsToDelete, + currentChain, + null + ); + + // Verify that the newChain is completed without any exceptions + assertNotNull(newChain); + assertTrue(newChain.isDone()); + assertFalse(newChain.isCompletedExceptionally()); + + // Verify that the deleteObjects method of S3AsyncClient was called with the expected request + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteObjectsRequest.class); + verify(s3AsyncClient, times(2)).deleteObjects(requestCaptor.capture()); + + List capturedRequests = requestCaptor.getAllValues(); + assertEquals(2, capturedRequests.size()); + + // Verify that the requests have the expected metric publisher added + for (DeleteObjectsRequest request : capturedRequests) { + assertNotNull(request.overrideConfiguration()); + assertTrue(request.overrideConfiguration().get().metricPublishers().contains(mockDeleteObjectsMetricPublisher)); + } + } + + public void testCreateDeleteBatches() { + List keys = Arrays.asList("key1", "key2", "key3", "key4", "key5", "key6"); + int bulkDeleteSize = 3; + + List> batches = S3AsyncDeleteHelper.createDeleteBatches(keys, bulkDeleteSize); + + assertEquals(2, batches.size()); + assertEquals(Arrays.asList("key1", "key2", "key3"), batches.get(0)); + assertEquals(Arrays.asList("key4", "key5", "key6"), batches.get(1)); + } + + public void testExecuteSingleDeleteBatch() throws Exception { + List batch = Arrays.asList("key1", "key2"); + DeleteObjectsResponse expectedResponse = DeleteObjectsResponse.builder().build(); + + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn(CompletableFuture.completedFuture(expectedResponse)); + + // Mock the getStatsMetricPublisher method of S3BlobStore to return a non-null value + StatsMetricPublisher mockMetricPublisher = mock(StatsMetricPublisher.class); + MetricPublisher mockDeleteObjectsMetricPublisher = mock(MetricPublisher.class); + when(blobStore.getStatsMetricPublisher()).thenReturn(mockMetricPublisher); + when(mockMetricPublisher.getDeleteObjectsMetricPublisher()).thenReturn(mockDeleteObjectsMetricPublisher); + + CompletableFuture future = S3AsyncDeleteHelper.executeSingleDeleteBatch(s3AsyncClient, blobStore, batch); + + assertNotNull(future); + assertTrue(future.isDone()); + assertFalse(future.isCompletedExceptionally()); + future.join(); // Wait for the CompletableFuture to complete + + // Verify that the deleteObjects method of S3AsyncClient was called with the expected request + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteObjectsRequest.class); + verify(s3AsyncClient).deleteObjects(requestCaptor.capture()); + + DeleteObjectsRequest capturedRequest = requestCaptor.getValue(); + assertEquals(blobStore.bucket(), capturedRequest.bucket()); + assertEquals(batch.size(), capturedRequest.delete().objects().size()); + assertTrue(capturedRequest.delete().objects().stream().map(ObjectIdentifier::key).collect(Collectors.toList()).containsAll(batch)); + } + + public void testProcessDeleteResponse() { + DeleteObjectsResponse response = DeleteObjectsResponse.builder() + .errors( + Arrays.asList( + S3Error.builder().key("key1").code("Code1").message("Message1").build(), + S3Error.builder().key("key2").code("Code2").message("Message2").build() + ) + ) + .build(); + + // Call the processDeleteResponse method + S3AsyncDeleteHelper.processDeleteResponse(response); + } + + public void testBulkDelete() { + List blobs = Arrays.asList("key1", "key2", "key3"); + String bucketName = "my-bucket"; + + // Mock the getStatsMetricPublisher method of S3BlobStore to return a non-null value + StatsMetricPublisher mockMetricPublisher = mock(StatsMetricPublisher.class); + MetricPublisher mockDeleteObjectsMetricPublisher = mock(MetricPublisher.class); + when(blobStore.getStatsMetricPublisher()).thenReturn(mockMetricPublisher); + when(mockMetricPublisher.getDeleteObjectsMetricPublisher()).thenReturn(mockDeleteObjectsMetricPublisher); + + DeleteObjectsRequest request = S3AsyncDeleteHelper.bulkDelete(bucketName, blobs, blobStore); + + assertEquals(bucketName, request.bucket()); + assertEquals(blobs.size(), request.delete().objects().size()); + assertTrue(request.delete().objects().stream().map(ObjectIdentifier::key).collect(Collectors.toList()).containsAll(blobs)); + assertNotNull(request.overrideConfiguration()); + assertTrue(request.overrideConfiguration().get().metricPublishers().contains(mockDeleteObjectsMetricPublisher)); + } + + public void testExecuteDeleteBatches() { + List> batches = Arrays.asList(Arrays.asList("key1", "key2"), Arrays.asList("key3", "key4")); + DeleteObjectsResponse expectedResponse = DeleteObjectsResponse.builder().build(); + + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn(CompletableFuture.completedFuture(expectedResponse)); + + // Mock the getStatsMetricPublisher method of S3BlobStore to return a non-null value + StatsMetricPublisher mockMetricPublisher = mock(StatsMetricPublisher.class); + MetricPublisher mockDeleteObjectsMetricPublisher = mock(MetricPublisher.class); + when(blobStore.getStatsMetricPublisher()).thenReturn(mockMetricPublisher); + when(mockMetricPublisher.getDeleteObjectsMetricPublisher()).thenReturn(mockDeleteObjectsMetricPublisher); + + CompletableFuture future = S3AsyncDeleteHelper.executeDeleteBatches(s3AsyncClient, blobStore, batches); + + assertNotNull(future); + assertTrue(future.isDone()); + assertFalse(future.isCompletedExceptionally()); + future.join(); // Wait for the CompletableFuture to complete + + // Verify that the deleteObjects method of S3AsyncClient was called with the expected requests + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteObjectsRequest.class); + verify(s3AsyncClient, times(2)).deleteObjects(requestCaptor.capture()); + + List capturedRequests = requestCaptor.getAllValues(); + assertEquals(2, capturedRequests.size()); + for (DeleteObjectsRequest request : capturedRequests) { + assertNotNull(request.overrideConfiguration()); + assertTrue(request.overrideConfiguration().get().metricPublishers().contains(mockDeleteObjectsMetricPublisher)); + } + } + + public void testExecuteDeleteChainWithAfterDeleteAction() { + List objectsToDelete = Arrays.asList("key1", "key2", "key3"); + CompletableFuture currentChain = CompletableFuture.completedFuture(null); + Runnable afterDeleteAction = mock(Runnable.class); + + // Mock the deleteObjects method of S3AsyncClient + when(s3AsyncClient.deleteObjects(any(DeleteObjectsRequest.class))).thenReturn( + CompletableFuture.completedFuture(DeleteObjectsResponse.builder().build()) + ); + + // Mock the getBulkDeletesSize method of S3BlobStore + when(blobStore.getBulkDeletesSize()).thenReturn(2); + + // Mock the getStatsMetricPublisher method of S3BlobStore to return a non-null value + StatsMetricPublisher mockMetricPublisher = mock(StatsMetricPublisher.class); + MetricPublisher mockDeleteObjectsMetricPublisher = mock(MetricPublisher.class); + when(blobStore.getStatsMetricPublisher()).thenReturn(mockMetricPublisher); + when(mockMetricPublisher.getDeleteObjectsMetricPublisher()).thenReturn(mockDeleteObjectsMetricPublisher); + + CompletableFuture newChain = S3AsyncDeleteHelper.executeDeleteChain( + s3AsyncClient, + blobStore, + objectsToDelete, + currentChain, + afterDeleteAction + ); + + // Verify that the newChain is completed without any exceptions + assertNotNull(newChain); + assertTrue(newChain.isDone()); + assertFalse(newChain.isCompletedExceptionally()); + + // Verify that the afterDeleteAction was called + verify(afterDeleteAction).run(); + } + +}