From 0175d3d358bcdf1a5e39bc862f383ed3d44d6135 Mon Sep 17 00:00:00 2001
From: Hui Zhao <zhohz@amazon.com>
Date: Thu, 29 Jun 2023 12:32:48 -0700
Subject: [PATCH] fix(storage): multipart upload cannot complete - When all
 parts have previously uploaded successfully

---
 .../providers/AWSS3UploadTask-unit-test.ts    | 73 +++++++++++++++++++
 .../storage/src/providers/AWSS3UploadTask.ts  |  6 +-
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/packages/storage/__tests__/providers/AWSS3UploadTask-unit-test.ts b/packages/storage/__tests__/providers/AWSS3UploadTask-unit-test.ts
index 11dcc206a65..52780efd31d 100644
--- a/packages/storage/__tests__/providers/AWSS3UploadTask-unit-test.ts
+++ b/packages/storage/__tests__/providers/AWSS3UploadTask-unit-test.ts
@@ -240,6 +240,79 @@ describe('resumable upload task test', () => {
 		);
 	});
 
+	test('Should complete upload if all parts have been uploaded on resume', done => {
+		const file = new File(['TestFileContent'], 'testFileName');
+		Object.defineProperty(file, 'size', { value: 25048576 });
+		const emitter = new events.EventEmitter();
+		const input: AWSS3UploadTaskParams = {
+			file,
+			// s3Client: new S3Client(testOpts),
+			s3Config: defaultS3Config,
+			emitter: emitter,
+			storage: mockLocalStorage,
+			level: 'public',
+			params: {
+				Bucket: 'bucket',
+				Key: 'key',
+			},
+			prefixPromise: Promise.resolve('prefix'),
+		};
+		(listParts as jest.Mock).mockResolvedValue({
+			Parts: [
+				{
+					PartNumber: 1,
+					Size: 25048576 / 2,
+					ETag: 'etag-1',
+				},
+				{
+					PartNumber: 2,
+					Size: 25048576 / 2,
+					ETag: 'etag-2',
+				},
+			],
+		});
+		(createMultipartUpload as jest.Mock).mockResolvedValue({
+			UploadId: 'uploadId',
+		});
+		(listObjectsV2 as jest.Mock).mockResolvedValue({
+			Contents: [{ Key: 'prefix' + input.params.Key, Size: 25048576 }],
+		});
+		(completeMultipartUpload as jest.Mock).mockResolvedValue({
+			Key: input.params.Key,
+		});
+		const fileMetadata: FileMetadata = {
+			bucket: 'bucket',
+			key: 'key',
+			lastTouched: Date.now(),
+			uploadId: 'uploadId',
+			fileName: file.name,
+		};
+		const fileId = [
+			file.name,
+			file.lastModified,
+			file.size,
+			file.type,
+			input.params.Bucket,
+			input.level,
+			input.params.Key,
+		].join('-');
+		const cachedUploadTasks = {
+			[fileId]: fileMetadata,
+		};
+		mockLocalStorage.setItem(
+			UPLOADS_STORAGE_KEY,
+			JSON.stringify(cachedUploadTasks)
+		);
+		const uploadTask = new AWSS3UploadTask(input);
+		// kick off the upload task
+		uploadTask.resume();
+
+		emitter.on(TaskEvents.UPLOAD_COMPLETE, () => {
+			expect(completeMultipartUpload).toBeCalledTimes(1);
+			done();
+		});
+	});
+
 	test('upload a body that exceeds the size of default part size and parts count', done => {
 		const testUploadId = 'testUploadId';
 		let buffer: ArrayBuffer;
diff --git a/packages/storage/src/providers/AWSS3UploadTask.ts b/packages/storage/src/providers/AWSS3UploadTask.ts
index 0853d248245..5e81c6a5846 100644
--- a/packages/storage/src/providers/AWSS3UploadTask.ts
+++ b/packages/storage/src/providers/AWSS3UploadTask.ts
@@ -447,7 +447,11 @@ export class AWSS3UploadTask implements UploadTask {
 				this.uploadId = uploadId;
 				this.queued = this._createParts();
 				this._initCachedUploadParts(parts);
-				this._startUpload();
+				if (this._isDone()) {
+					this._completeUpload();
+				} else {
+					this._startUpload();
+				}
 			} else {
 				if (!this.uploadId) {
 					const uploadId = await this._initMultipartUpload();