Skip to content

Commit

Permalink
chore: add tests for TransferManager
Browse files Browse the repository at this point in the history
  • Loading branch information
@kessplas
kessplas committed Oct 30, 2024
1 parent 37e772f commit 684bc47
Show file tree
Hide file tree
Showing 2 changed files with 152 additions and 0 deletions.
7 changes: 7 additions & 0 deletions pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,13 @@
<scope>test</scope>
</dependency>

<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>s3-transfer-manager</artifactId>
<optional>true</optional>
<version>2.28.28</version>
</dependency>

<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk-kms</artifactId>
Expand Down
145 changes: 145 additions & 0 deletions src/test/java/software/amazon/encryption/s3/S3AsyncEncryptionClientTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,19 @@
import software.amazon.awssdk.services.s3.model.CopyObjectResponse;
import software.amazon.awssdk.services.s3.model.DeleteObjectResponse;
import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse;
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
import software.amazon.awssdk.services.s3.model.ObjectIdentifier;
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
import software.amazon.awssdk.services.s3.model.S3Exception;
import software.amazon.awssdk.services.s3.multipart.MultipartConfiguration;
import software.amazon.awssdk.transfer.s3.S3TransferManager;
import software.amazon.awssdk.transfer.s3.model.CompletedDownload;
import software.amazon.awssdk.transfer.s3.model.Download;
import software.amazon.awssdk.transfer.s3.model.DownloadRequest;
import software.amazon.awssdk.transfer.s3.model.Upload;
import software.amazon.awssdk.transfer.s3.model.UploadRequest;
import software.amazon.awssdk.transfer.s3.progress.LoggingTransferListener;
import software.amazon.encryption.s3.internal.InstructionFileConfig;
import software.amazon.encryption.s3.materials.KmsKeyring;
import software.amazon.encryption.s3.utils.BoundedInputStream;
Expand All @@ -61,6 +69,7 @@
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

Expand Down Expand Up @@ -141,6 +150,142 @@ public void asyncCustomConfiguration() {
s3Client.close();
}

@Test
public void transferManagerUploadString() {
final String objectKey = appendTestSuffix("tm-string");
final String input = "short test of s3 encryption client with transfer manager";
S3AsyncClient v3AsyncClient = S3AsyncEncryptionClient.builder()
.aesKey(AES_KEY)
.build();
S3TransferManager transferManager =
S3TransferManager.builder()
.s3Client(v3AsyncClient)
.build();

Upload upload = transferManager.upload(UploadRequest.builder()
.putObjectRequest((builder -> builder
.bucket(BUCKET)
.key(objectKey)
.build()))
.requestBody(AsyncRequestBody.fromString(input))
.build());
upload.completionFuture().join();

// tm download
Download<ResponseBytes<GetObjectResponse>> download = transferManager.download(DownloadRequest.builder()
.getObjectRequest(GetObjectRequest.builder()
.bucket(BUCKET)
.key(objectKey)
.build())
.responseTransformer(AsyncResponseTransformer.toBytes())
.build());
CompletedDownload<ResponseBytes<GetObjectResponse>> resp = download.completionFuture().join();
assertEquals(input, resp.result().asUtf8String());

// Cleanup
deleteObject(BUCKET, objectKey, v3AsyncClient);
transferManager.close();
}

@Test
public void transferManagerUploadStream() throws IOException {
final String objectKey = appendTestSuffix("tm-stream");

final long fileSizeLimit = 1024 * 1024 * 100;
final InputStream inputStream = new BoundedInputStream(fileSizeLimit);
final InputStream objectStreamForResult = new BoundedInputStream(fileSizeLimit);
final InputStream objectStreamForResultTm = new BoundedInputStream(fileSizeLimit);

S3AsyncClient v3AsyncClient = S3AsyncEncryptionClient.builder()
.aesKey(AES_KEY)
.enableDelayedAuthenticationMode(true)
.enableMultipartPutObject(true)
.build();
S3TransferManager transferManager =
S3TransferManager.builder()
.s3Client(v3AsyncClient)
.build();

ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
Upload upload = transferManager.upload(UploadRequest.builder()
.putObjectRequest((builder -> builder
.bucket(BUCKET)
.key(objectKey)
.build()))
.requestBody(AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, singleThreadExecutor))
.addTransferListener(LoggingTransferListener.create())
.build());
upload.completionFuture().join();
singleThreadExecutor.shutdown();

// tm download
Download<ResponseInputStream<GetObjectResponse>> download = transferManager.download(DownloadRequest.builder()
.getObjectRequest(GetObjectRequest.builder()
.bucket(BUCKET)
.key(objectKey)
.build())
.responseTransformer(AsyncResponseTransformer.toBlockingInputStream())
.build());

CompletedDownload<ResponseInputStream<GetObjectResponse>> resp = download.completionFuture().join();
assertTrue(IOUtils.contentEquals(objectStreamForResultTm, resp.result()));

// Cleanup
deleteObject(BUCKET, objectKey, v3AsyncClient);
transferManager.close();
}

@Test
public void transferManagerUploadStreamCrt() throws ExecutionException, InterruptedException, IOException {
final String objectKey = appendTestSuffix("tm-stream-crt");

final long fileSizeLimit = 1024 * 1024 * 100;
final InputStream inputStream = new BoundedInputStream(fileSizeLimit);
final InputStream objectStreamForResult = new BoundedInputStream(fileSizeLimit);
final InputStream objectStreamForResultTm = new BoundedInputStream(fileSizeLimit);

S3AsyncClient wrappedCrt = S3AsyncClient.crtBuilder()
.minimumPartSizeInBytes(8000000L)
.thresholdInBytes(500L)
.build();
S3AsyncClient v3AsyncClient = S3AsyncEncryptionClient.builder()
.wrappedClient(wrappedCrt)
.aesKey(AES_KEY)
.enableDelayedAuthenticationMode(true)
.enableMultipartPutObject(true)
.build();
S3TransferManager transferManager =
S3TransferManager.builder()
.s3Client(v3AsyncClient)
.build();

ExecutorService singleThreadExecutor = Executors.newSingleThreadExecutor();
Upload upload = transferManager.upload(UploadRequest.builder()
.putObjectRequest((builder -> builder
.bucket(BUCKET)
.key(objectKey)
.build()))
.requestBody(AsyncRequestBody.fromInputStream(inputStream, fileSizeLimit, singleThreadExecutor))
.addTransferListener(LoggingTransferListener.create())
.build());
upload.completionFuture().join();
singleThreadExecutor.shutdown();

Download<ResponseInputStream<GetObjectResponse>> download = transferManager.download(DownloadRequest.builder()
.getObjectRequest(GetObjectRequest.builder()
.bucket(BUCKET)
.key(objectKey)
.build())
.responseTransformer(AsyncResponseTransformer.toBlockingInputStream())
.build());
download.completionFuture().join();
CompletedDownload<ResponseInputStream<GetObjectResponse>> resp = download.completionFuture().get();

assertTrue(IOUtils.contentEquals(objectStreamForResultTm, resp.result()));
deleteObject(BUCKET, objectKey, v3AsyncClient);
transferManager.close();
}

@Test
public void asyncTopLevelConfigurationAllOptions() {
final String objectKey = appendTestSuffix("async-top-level-all-options");
Expand Down

0 comments on commit 684bc47

Please sign in to comment.