Feat: DownloadData refactor (#13073)

* updating type info * update downloadData implementation * fix output type * path dynamic construction * resolve merge conflit * add unit tests for path * update bundle size * update deprecation warning * update test * misc changes * address comments
aws-amplify · Mar 11, 2024 · 32990ee · 32990ee
1 parent f049a9e
commit 32990ee
Show file tree

Hide file tree

Showing 19 changed files with 465 additions and 93 deletions.
diff --git a/packages/aws-amplify/package.json b/packages/aws-amplify/package.json
@@ -466,7 +466,7 @@
 			"name": "[Storage] downloadData (S3)",
 			"path": "./dist/esm/storage/index.mjs",
 			"import": "{ downloadData }",
-			"limit": "14.00 kB"
+			"limit": "14.10 kB"
 		},
 		{
 			"name": "[Storage] getProperties (S3)",

diff --git a/packages/storage/__tests__/providers/s3/apis/downloadData.test.ts b/packages/storage/__tests__/providers/s3/apis/downloadData.test.ts
@@ -5,8 +5,18 @@ import { AWSCredentials } from '@aws-amplify/core/internals/utils';
 import { Amplify } from '@aws-amplify/core';
 import { getObject } from '../../../../src/providers/s3/utils/client';
 import { downloadData } from '../../../../src/providers/s3';
-import { createDownloadTask } from '../../../../src/providers/s3/utils';
-import { DownloadDataOptions } from '../../../../src/providers/s3/types';
+import {
+	createDownloadTask,
+	validateStorageOperationInput,
+} from '../../../../src/providers/s3/utils';
+import {
+	DownloadDataOptionsKey,
+	DownloadDataOptionsPath,
+} from '../../../../src/providers/s3/types';
+import {
+	STORAGE_INPUT_KEY,
+	STORAGE_INPUT_PATH,
+} from '../../../../src/providers/s3/utils/constants';
 
 jest.mock('../../../../src/providers/s3/utils/client');
 jest.mock('../../../../src/providers/s3/utils');
@@ -34,9 +44,10 @@ const defaultIdentityId = 'defaultIdentityId';
 
 const mockFetchAuthSession = Amplify.Auth.fetchAuthSession as jest.Mock;
 const mockCreateDownloadTask = createDownloadTask as jest.Mock;
+const mockValidateStorageInput = validateStorageOperationInput as jest.Mock;
 const mockGetConfig = Amplify.getConfig as jest.Mock;
 
-describe('downloadData', () => {
+describe('downloadData with key', () => {
 	beforeAll(() => {
 		mockFetchAuthSession.mockResolvedValue({
 			credentials,
@@ -52,16 +63,20 @@ describe('downloadData', () => {
 		});
 	});
 	mockCreateDownloadTask.mockReturnValue('downloadTask');
+	mockValidateStorageInput.mockReturnValue({
+		inputType: STORAGE_INPUT_KEY,
+		objectKey: key,
+	});
 
 	beforeEach(() => {
 		jest.clearAllMocks();
 	});
 
-	it('should return a download task', async () => {
+	it('should return a download task with key', async () => {
 		expect(downloadData({ key: 'key' })).toBe('downloadTask');
 	});
 
-	[
+	test.each([
 		{
 			expectedKey: `public/${key}`,
 		},
@@ -81,14 +96,9 @@ describe('downloadData', () => {
 			options: { accessLevel: 'protected', targetIdentityId },
 			expectedKey: `protected/${targetIdentityId}/${key}`,
 		},
-	].forEach(({ options, expectedKey }) => {
-		const accessLevelMsg = options?.accessLevel ?? 'default';
-		const targetIdentityIdMsg = options?.targetIdentityId
-			? `and targetIdentityId`
-			: '';
-
-		it(`should supply the correct parameters to getObject API handler with ${accessLevelMsg} accessLevel ${targetIdentityIdMsg}`, async () => {
-			expect.assertions(2);
+	])(
+		'should supply the correct parameters to getObject API handler with $expectedKey accessLevel',
+		async ({ options, expectedKey }) => {
 			(getObject as jest.Mock).mockResolvedValueOnce({ Body: 'body' });
 			const onProgress = jest.fn();
 			downloadData({
@@ -97,7 +107,7 @@ describe('downloadData', () => {
 					...options,
 					useAccelerateEndpoint: true,
 					onProgress,
-				} as DownloadDataOptions,
+				} as DownloadDataOptionsKey,
 			});
 			const job = mockCreateDownloadTask.mock.calls[0][0].job;
 			await job();
@@ -116,11 +126,10 @@ describe('downloadData', () => {
 					Key: expectedKey,
 				},
 			);
-		});
-	});
+		},
+	);
 
-	it('should assign the getObject API handler response to the result', async () => {
-		expect.assertions(2);
+	it('should assign the getObject API handler response to the result with key', async () => {
 		const lastModified = 'lastModified';
 		const contentLength = 'contentLength';
 		const eTag = 'eTag';
@@ -177,3 +186,131 @@ describe('downloadData', () => {
 		);
 	});
 });
+
+describe('downloadData with path', () => {
+	beforeAll(() => {
+		mockFetchAuthSession.mockResolvedValue({
+			credentials,
+			identityId: defaultIdentityId,
+		});
+		mockGetConfig.mockReturnValue({
+			Storage: {
+				S3: {
+					bucket,
+					region,
+				},
+			},
+		});
+		mockCreateDownloadTask.mockReturnValue('downloadTask');
+		mockValidateStorageInput.mockReturnValue({
+			inputType: STORAGE_INPUT_PATH,
+			objectKey: 'path',
+		});
+	});
+
+	beforeEach(() => {
+		jest.clearAllMocks();
+	});
+
+	it('should return a download task with path', async () => {
+		expect(downloadData({ path: 'path' })).toBe('downloadTask');
+	});
+
+	test.each([
+		{
+			path: 'path',
+			expectedKey: 'path',
+		},
+		{
+			path: () => 'path',
+			expectedKey: 'path',
+		},
+	])(
+		'should call getObject API with $expectedKey when path provided is $path',
+		async ({ path, expectedKey }) => {
+			(getObject as jest.Mock).mockResolvedValueOnce({ Body: 'body' });
+			const onProgress = jest.fn();
+			downloadData({
+				path: path,
+				options: {
+					useAccelerateEndpoint: true,
+					onProgress,
+				} as DownloadDataOptionsPath,
+			});
+			const job = mockCreateDownloadTask.mock.calls[0][0].job;
+			await job();
+			expect(getObject).toHaveBeenCalledTimes(1);
+			expect(getObject).toHaveBeenCalledWith(
+				{
+					credentials,
+					region,
+					useAccelerateEndpoint: true,
+					onDownloadProgress: onProgress,
+					abortSignal: expect.any(AbortSignal),
+					userAgentValue: expect.any(String),
+				},
+				{
+					Bucket: bucket,
+					Key: expectedKey,
+				},
+			);
+		},
+	);
+
+	it('should assign the getObject API handler response to the result with path', async () => {
+		const lastModified = 'lastModified';
+		const contentLength = 'contentLength';
+		const eTag = 'eTag';
+		const metadata = 'metadata';
+		const versionId = 'versionId';
+		const contentType = 'contentType';
+		const body = 'body';
+		const path = 'path';
+		(getObject as jest.Mock).mockResolvedValueOnce({
+			Body: body,
+			LastModified: lastModified,
+			ContentLength: contentLength,
+			ETag: eTag,
+			Metadata: metadata,
+			VersionId: versionId,
+			ContentType: contentType,
+		});
+		downloadData({ path });
+		const job = mockCreateDownloadTask.mock.calls[0][0].job;
+		const result = await job();
+		expect(getObject).toHaveBeenCalledTimes(1);
+		expect(result).toEqual({
+			path,
+			body,
+			lastModified,
+			size: contentLength,
+			eTag,
+			metadata,
+			versionId,
+			contentType,
+		});
+	});
+
+	it('should forward the bytes range option to the getObject API', async () => {
+		const start = 1;
+		const end = 100;
+		(getObject as jest.Mock).mockResolvedValueOnce({ Body: 'body' });
+
+		downloadData({
+			path: 'mockPath',
+			options: {
+				bytesRange: { start, end },
+			},
+		});
+
+		const job = mockCreateDownloadTask.mock.calls[0][0].job;
+		await job();
+
+		expect(getObject).toHaveBeenCalledWith(
+			expect.anything(),
+			expect.objectContaining({
+				Range: `bytes=${start}-${end}`,
+			}),
+		);
+	});
+});
diff --git a/packages/storage/__tests__/providers/s3/apis/utils/validateStorageOperationInput.test.ts b/packages/storage/__tests__/providers/s3/apis/utils/validateStorageOperationInput.test.ts
@@ -0,0 +1,53 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+	StorageValidationErrorCode,
+	validationErrorMap,
+} from '../../../../../src/errors/types/validation';
+import { validateStorageOperationInput } from '../../../../../src/providers/s3/utils';
+import {
+	STORAGE_INPUT_KEY,
+	STORAGE_INPUT_PATH,
+} from '../../../../../src/providers/s3/utils/constants';
+
+describe('validateStorageOperationInput', () => {
+	it('should return inputType as STORAGE_INPUT_PATH and objectKey as testPath when input is path as string', () => {
+		const input = { path: 'testPath' };
+		const result = validateStorageOperationInput(input);
+		expect(result).toEqual({
+			inputType: STORAGE_INPUT_PATH,
+			objectKey: 'testPath',
+		});
+	});
+
+	it('should return inputType as STORAGE_INPUT_PATH and objectKey as result of path function when input is path as function', () => {
+		const input = {
+			path: ({ identityId }: { identityId?: string }) =>
+				`testPath/${identityId}`,
+		};
+		const result = validateStorageOperationInput(input, '123');
+		expect(result).toEqual({
+			inputType: STORAGE_INPUT_PATH,
+			objectKey: 'testPath/123',
+		});
+	});
+
+	it('should return inputType as STORAGE_INPUT_KEY and objectKey as testKey when input is key', () => {
+		const input = { key: 'testKey' };
+		const result = validateStorageOperationInput(input);
+		expect(result).toEqual({
+			inputType: STORAGE_INPUT_KEY,
+			objectKey: 'testKey',
+		});
+	});
+
+	it('should throw an error when input is invalid', () => {
+		const input = { invalid: 'test' } as any;
+		expect(() => validateStorageOperationInput(input)).toThrow(
+			validationErrorMap[
+				StorageValidationErrorCode.InvalidStorageOperationInput
+			].message,
+		);
+	});
+});
diff --git a/packages/storage/src/errors/types/validation.ts b/packages/storage/src/errors/types/validation.ts
@@ -14,6 +14,7 @@ export enum StorageValidationErrorCode {
 	UrlExpirationMaxLimitExceed = 'UrlExpirationMaxLimitExceed',
 	ObjectIsTooLarge = 'ObjectIsTooLarge',
 	InvalidUploadSource = 'InvalidUploadSource',
+	InvalidStorageOperationInput = 'InvalidStorageOperationInput',
 }
 
 export const validationErrorMap: AmplifyErrorMap<StorageValidationErrorCode> = {
@@ -49,4 +50,7 @@ export const validationErrorMap: AmplifyErrorMap<StorageValidationErrorCode> = {
 		message:
 			'Upload source type can only be a `Blob`, `File`, `ArrayBuffer`, or `string`.',
 	},
+	[StorageValidationErrorCode.InvalidStorageOperationInput]: {
+		message: 'Missing path or key parameter in Input',
+	},
 };