Skip to content

Commit

Permalink
Adding UTs for S3BlobContainer
Browse files Browse the repository at this point in the history 
Signed-off-by: Raghuvansh Raj <raghraaj@amazon.com>
  • Loading branch information
@raghuvanshraj
raghuvanshraj committed Jun 2, 2023
1 parent c855e61 commit 5ceface
Show file tree
Hide file tree
Showing 2 changed files with 246 additions and 1 deletion.
2 changes: 1 addition & 1 deletion plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -246,7 +246,7 @@ private void doDeleteBlobs(List<String> blobNames, boolean relative) throws IOEx
deleteRequests.add(bulkDelete(blobStore.bucket(), partition));
}
SocketAccess.doPrivilegedVoid(() -> {
RuntimeException aex = null;
SdkException aex = null;
for (DeleteObjectsRequest deleteRequest : deleteRequests) {
List<String> keysInRequest = deleteRequest.delete()
.objects()
Expand Down
245 changes: 245 additions & 0 deletions ...repository-s3/src/test/java/org/opensearch/repositories/s3/S3BlobStoreContainerTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import org.mockito.ArgumentCaptor;
import org.opensearch.common.blobstore.BlobPath;
import org.opensearch.common.blobstore.BlobStoreException;
import org.opensearch.common.blobstore.DeleteResult;
import org.opensearch.common.collect.Tuple;
import org.opensearch.common.unit.ByteSizeUnit;
import org.opensearch.test.OpenSearchTestCase;
Expand All @@ -48,20 +49,36 @@
import software.amazon.awssdk.services.s3.model.CompletedPart;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest;
import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse;
import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;
import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
import software.amazon.awssdk.services.s3.model.ObjectIdentifier;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
import software.amazon.awssdk.services.s3.model.S3Error;
import software.amazon.awssdk.services.s3.model.S3Object;
import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
import software.amazon.awssdk.services.s3.model.StorageClass;
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 java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

Expand Down Expand Up @@ -102,6 +119,234 @@ public void testExecuteSingleUploadBlobSizeLargerThanBufferSize() {
assertEquals("Upload request size [2097152] can't be larger than buffer size", e.getMessage());
}

public void testBlobExists() {
final String bucketName = randomAlphaOfLengthBetween(1, 10);
final String blobName = randomAlphaOfLengthBetween(1, 10);

final BlobPath blobPath = new BlobPath();

final S3BlobStore blobStore = mock(S3BlobStore.class);
when(blobStore.bucket()).thenReturn(bucketName);

final S3Client client = mock(S3Client.class);
when(client.headObject(any(HeadObjectRequest.class))).thenReturn(HeadObjectResponse.builder().build());
final AmazonS3Reference clientReference = new AmazonS3Reference(client);
when(blobStore.clientReference()).thenReturn(clientReference);

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

assertTrue(blobContainer.blobExists(blobName));
verify(client, times(1)).headObject(any(HeadObjectRequest.class));
}

public void testBlobExistsNoSuchKeyException() {
final String bucketName = randomAlphaOfLengthBetween(1, 10);
final String blobName = randomAlphaOfLengthBetween(1, 10);

final BlobPath blobPath = new BlobPath();

final S3BlobStore blobStore = mock(S3BlobStore.class);
when(blobStore.bucket()).thenReturn(bucketName);

final S3Client client = mock(S3Client.class);
when(client.headObject(any(HeadObjectRequest.class))).thenThrow(NoSuchKeyException.builder().build());
final AmazonS3Reference clientReference = new AmazonS3Reference(client);
when(blobStore.clientReference()).thenReturn(clientReference);

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

assertFalse(blobContainer.blobExists(blobName));
verify(client, times(1)).headObject(any(HeadObjectRequest.class));
}

public void testBlobExistsRequestFailure() {
final String bucketName = randomAlphaOfLengthBetween(1, 10);
final String blobName = randomAlphaOfLengthBetween(1, 10);

final BlobPath blobPath = new BlobPath();

final S3BlobStore blobStore = mock(S3BlobStore.class);
when(blobStore.bucket()).thenReturn(bucketName);

final S3Client client = mock(S3Client.class);
final AmazonS3Reference clientReference = new AmazonS3Reference(client);
when(blobStore.clientReference()).thenReturn(clientReference);

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

when(client.headObject(any(HeadObjectRequest.class))).thenThrow(new RuntimeException());

assertThrows(BlobStoreException.class, () -> blobContainer.blobExists(blobName));
verify(client, times(1)).headObject(any(HeadObjectRequest.class));
}

public void testDelete() throws IOException {
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());

final S3Client client = mock(S3Client.class);
doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference();

ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class);
Iterator<ListObjectsV2Response> listObjectsV2ResponseIterator = (Iterator<ListObjectsV2Response>) mock(Iterator.class);
when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator);
final AtomicInteger iteratorNextInvocationCount = new AtomicInteger();
final int listS3ObjectsPageCount = 3;
final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5);
final int s3ObjectsPerPage = 5;
final List<String> keysListed = new ArrayList<>();
doAnswer(invocation -> iteratorNextInvocationCount.get() < listS3ObjectsPageCount).when(listObjectsV2ResponseIterator).hasNext();
doAnswer(invocation -> {
if (iteratorNextInvocationCount.getAndIncrement() < listS3ObjectsPageCount) {
String s3ObjectKey = UUID.randomUUID().toString();
keysListed.add(s3ObjectKey);
return ListObjectsV2Response.builder()
.contents(Collections.nCopies(s3ObjectsPerPage, S3Object.builder().key(s3ObjectKey).size(s3ObjectSize).build()))
.build();
}
throw new NoSuchElementException();
}).when(listObjectsV2ResponseIterator).next();
when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable);

final List<String> keysDeleted = new ArrayList<>();
doAnswer(invocation -> {
DeleteObjectsRequest deleteObjectsRequest = invocation.getArgument(0);
keysDeleted.addAll(deleteObjectsRequest.delete().objects().stream().map(ObjectIdentifier::key).collect(Collectors.toList()));
return DeleteObjectsResponse.builder().build();
}).when(client).deleteObjects(any(DeleteObjectsRequest.class));

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

DeleteResult deleteResult = blobContainer.delete();
assertEquals(s3ObjectSize * s3ObjectsPerPage * listS3ObjectsPageCount, deleteResult.bytesDeleted());
assertEquals(s3ObjectsPerPage * listS3ObjectsPageCount, deleteResult.blobsDeleted());
// keysDeleted will have blobPath also
assertEquals(keysListed.size(), keysDeleted.size() - 1);
assertTrue(keysDeleted.contains(blobPath.buildAsString()));
assertArrayEquals(
keysListed.toArray(String[]::new),
keysDeleted.stream().filter(key -> !blobPath.buildAsString().equals(key)).toArray(String[]::new)
);
}

public void testDeleteItemLevelErrorsDuringDelete() throws IOException {
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());

final S3Client client = mock(S3Client.class);
doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference();

ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class);
Iterator<ListObjectsV2Response> listObjectsV2ResponseIterator = (Iterator<ListObjectsV2Response>) mock(Iterator.class);
when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator);
final AtomicInteger iteratorNextInvocationCount = new AtomicInteger();
final int listS3ObjectsPageCount = 3;
final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5);
final int s3ObjectsPerPage = 5;
doAnswer(invocation -> iteratorNextInvocationCount.get() < listS3ObjectsPageCount).when(listObjectsV2ResponseIterator).hasNext();
doAnswer(invocation -> {
if (iteratorNextInvocationCount.getAndIncrement() < listS3ObjectsPageCount) {
String s3ObjectKey = UUID.randomUUID().toString();
return ListObjectsV2Response.builder()
.contents(Collections.nCopies(s3ObjectsPerPage, S3Object.builder().key(s3ObjectKey).size(s3ObjectSize).build()))
.build();
}
throw new NoSuchElementException();
}).when(listObjectsV2ResponseIterator).next();
when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable);

final List<String> keysFailedDeletion = new ArrayList<>();
doAnswer(invocation -> {
DeleteObjectsRequest deleteObjectsRequest = invocation.getArgument(0);
int i = 0;
for (ObjectIdentifier objectIdentifier : deleteObjectsRequest.delete().objects()) {
if (i % 2 == 0) {
keysFailedDeletion.add(objectIdentifier.key());
}
i++;
}
return DeleteObjectsResponse.builder()
.errors(keysFailedDeletion.stream().map(key -> S3Error.builder().key(key).build()).collect(Collectors.toList()))
.build();
}).when(client).deleteObjects(any(DeleteObjectsRequest.class));

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

assertThrows(AssertionError.class, blobContainer::delete);
}

public void testDeleteSdkExceptionDuringListOperation() throws IOException {
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());

final S3Client client = mock(S3Client.class);
doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference();

ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class);
Iterator<ListObjectsV2Response> listObjectsV2ResponseIterator = (Iterator<ListObjectsV2Response>) mock(Iterator.class);
when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator);
when(listObjectsV2ResponseIterator.hasNext()).thenReturn(true);
when(listObjectsV2ResponseIterator.next()).thenThrow(SdkException.builder().build());
when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable);

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

assertThrows(IOException.class, blobContainer::delete);
}

public void testDeleteSdkExceptionDuringDeleteOperation() throws IOException {
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());

final S3Client client = mock(S3Client.class);
doAnswer(invocation -> new AmazonS3Reference(client)).when(blobStore).clientReference();

ListObjectsV2Iterable listObjectsV2Iterable = mock(ListObjectsV2Iterable.class);
Iterator<ListObjectsV2Response> listObjectsV2ResponseIterator = (Iterator<ListObjectsV2Response>) mock(Iterator.class);
when(listObjectsV2Iterable.iterator()).thenReturn(listObjectsV2ResponseIterator);
final AtomicInteger iteratorNextInvocationCount = new AtomicInteger();
final int listS3ObjectsPageCount = 3;
final long s3ObjectSize = ByteSizeUnit.MB.toBytes(5);
final int s3ObjectsPerPage = 5;
doAnswer(invocation -> iteratorNextInvocationCount.get() < listS3ObjectsPageCount).when(listObjectsV2ResponseIterator).hasNext();
doAnswer(invocation -> {
if (iteratorNextInvocationCount.getAndIncrement() < listS3ObjectsPageCount) {
String s3ObjectKey = UUID.randomUUID().toString();
return ListObjectsV2Response.builder()
.contents(Collections.nCopies(s3ObjectsPerPage, S3Object.builder().key(s3ObjectKey).size(s3ObjectSize).build()))
.build();
}
throw new NoSuchElementException();
}).when(listObjectsV2ResponseIterator).next();
when(client.listObjectsV2Paginator(any(ListObjectsV2Request.class))).thenReturn(listObjectsV2Iterable);

when(client.deleteObjects(any(DeleteObjectsRequest.class))).thenThrow(SdkException.builder().build());

final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);

assertThrows(IOException.class, blobContainer::delete);
}

public void testExecuteSingleUpload() throws IOException {
final String bucketName = randomAlphaOfLengthBetween(1, 10);
final String blobName = randomAlphaOfLengthBetween(1, 10);
Expand Down

0 comments on commit 5ceface

Please sign in to comment.