StorageImage Component (#4054)

* fix: fix infinite loop in useStorageURL

* feat: add StorageImage component

* test: add unit tests for StorageImage

* chore: fix loading attribute

* chore: update snapshots

* test: update snapshots

* chore: address feedback

* chore: address comments

* chore: address comments

* test: fix test

* chore: address comments

* chore: remove unused comments

* chore: address comments

* chore: update tests

* chore: add comments to tests

docs/tests/__snapshots__/props-table.test.ts.snap

@@ -3768,6 +3768,13 @@ exports[`Props Table 1`] = `
                         \\"category\\": \\"ImageOptions\\",
                         \\"isOptional\\": true
                     },
+                    \\"aspectRatio\\": {
+                        \\"name\\": \\"aspectRatio\\",
+                        \\"type\\": \\"ResponsiveStyle<Property.AspectRatio> | undefined\\",
+                        \\"description\\": \\"Sets a preferred aspect ratio for the selected replaced element's box\\",
+                        \\"category\\": \\"ImageStyleProps\\",
+                        \\"isOptional\\": true
+                    },
                     \\"objectFit\\": {
                         \\"name\\": \\"objectFit\\",
                         \\"type\\": \\"ResponsiveStyle<Property.ObjectFit> | undefined\\",

packages/react-storage/src/components/StorageImage/StorageImage.tsx

@@ -0,0 +1,43 @@
+import * as React from 'react';
+import classNames from 'classnames';
+
+import { Image, ComponentClassNames } from '@aws-amplify/ui-react';
+import { useStorageURL } from '@aws-amplify/ui-react/internal';
+
+import type { StorageImageProps } from './types';
+
+export const StorageImage = ({
+  accessLevel,
+  className,
+  fallbackSrc,
+  identityId,
+  imgKey,
+  onStorageGetError,
+  ...rest
+}: StorageImageProps): JSX.Element => {
+  const options = React.useMemo(
+    () => ({
+      accessLevel,
+      identityId,
+    }),
+    [accessLevel, identityId]
+  );
+
+  const errorConfig = React.useMemo(
+    () => ({
+      fallbackURL: fallbackSrc,
+      onStorageGetError,
+    }),
+    [fallbackSrc, onStorageGetError]
+  );
+
+  const url = useStorageURL(imgKey, options, errorConfig);
+
+  return (
+    <Image
+      {...rest}
+      className={classNames(ComponentClassNames.StorageImage, className)}
+      src={url}
+    />
+  );
+};

packages/react-storage/src/components/StorageImage/_tests_/StorageImage.test.tsx

@@ -0,0 +1,78 @@
+import * as React from 'react';
+import { render, screen } from '@testing-library/react';
+
+import { Storage } from 'aws-amplify';
+import { ComponentClassNames } from '@aws-amplify/ui-react';
+
+import { StorageImage } from '../StorageImage';
+
+describe('StorageImage', () => {
+  const imgKey = 'test.jpg';
+  const imgURL = 'https://amplify.s3.amazonaws.com/path/to/test.jpg';
+  const fallbackSrc = 'https://amplify.s3.amazonaws.com/path/to/fallback.jpg';
+  const errorMessage = '500 Internal Server Error';
+
+  beforeAll(() => {
+    jest.spyOn(Storage, 'get').mockResolvedValue(imgURL);
+  });
+
+  it('should render default classname', async () => {
+    render(
+      <StorageImage alt="StorageImage" imgKey={imgKey} accessLevel="public" />
+    );
+
+    const img = await screen.findByRole('img');
+    expect(img).toHaveClass(ComponentClassNames.StorageImage);
+  });
+
+  it('should render custom classname', async () => {
+    const className = 'MyImage';
+    render(
+      <StorageImage
+        alt="StorageImage"
+        className={className}
+        imgKey={imgKey}
+        accessLevel="public"
+      />
+    );
+
+    const img = await screen.findByRole('img');
+    expect(img).toHaveClass(className);
+  });
+
+  it('should get the presigned URL and pass it to image src attribute', async () => {
+    const onStorageError = jest.fn();
+    render(
+      <StorageImage
+        alt="StorageImage"
+        imgKey={imgKey}
+        accessLevel="public"
+        onStorageGetError={onStorageError}
+      />
+    );
+
+    const img = await screen.findByRole('img');
+    expect(onStorageError).not.toHaveBeenCalled();
+    expect(img).toHaveAttribute('src', imgURL);
+  });
+
+  it('should set image src attribute to fallbackSrc and invoke onGetStorageError when Storage.get is rejected', async () => {
+    jest.restoreAllMocks();
+    jest.spyOn(Storage, 'get').mockRejectedValue(errorMessage);
+    const onStorageError = jest.fn();
+    render(
+      <StorageImage
+        alt="StorageImage"
+        imgKey={imgKey}
+        accessLevel="public"
+        fallbackSrc={fallbackSrc}
+        onStorageGetError={onStorageError}
+      />
+    );
+
+    const img = await screen.findByRole('img');
+    expect(onStorageError).toHaveBeenCalledTimes(1);
+    expect(onStorageError).toHaveBeenCalledWith(errorMessage);
+    expect(img).toHaveAttribute('src', fallbackSrc);
+  });
+});

packages/react-storage/src/components/StorageImage/index.ts

@@ -0,0 +1,2 @@
+export { StorageImage } from './StorageImage';
+export { StorageImageProps } from './types';

packages/react-storage/src/components/StorageImage/types.ts

@@ -0,0 +1,12 @@
+import type { StorageAccessLevel } from '@aws-amplify/storage';
+import type { ImageProps } from '@aws-amplify/ui-react';
+
+export interface StorageImageProps extends Omit<ImageProps, 'src'> {
+  // Use imgKey instead of key because key is a reserved keyword
+  // and cannot be accessed via props in React components
+  imgKey: string;
+  accessLevel: StorageAccessLevel;
+  identityId?: string;
+  fallbackSrc?: string;
+  onStorageGetError?: (error: Error) => void;
+}

packages/react-storage/src/components/index.ts

@@ -1,3 +1,5 @@
+export { StorageImage, StorageImageProps } from './StorageImage';
+
 export {
   StorageManager,
   StorageManagerProps,

packages/react-storage/src/index.ts

@@ -1,4 +1,6 @@
 export {
+  StorageImage,
+  StorageImageProps,
   StorageManager,
   StorageManagerProps,
   DropZoneProps,

packages/react/__tests__/__snapshots__/exports.ts.snap

@@ -5087,6 +5087,9 @@ Object {
       "ariaValuetext": Object {
         "type": "string",
       },
+      "aspectRatio": Object {
+        "type": "string",
+      },
       "backgroundColor": Object {
         "type": "string",
       },

packages/react/src/hooks/__tests__/useStorageURL.test.ts

@@ -1,4 +1,4 @@
-import { act, renderHook } from '@testing-library/react-hooks';
+import { renderHook } from '@testing-library/react-hooks';
 import { useStorageURL } from '../useStorageURL';
 import { S3ProviderGetConfig, Storage } from '@aws-amplify/storage';
 
@@ -10,82 +10,63 @@ describe('useStorageURL', () => {
   const storageKey = 'file.jpg';
   const storageOptions: S3ProviderGetConfig = { level: 'public' };
   const storageUrl = 'https://amplify.s3.amazonaws.com/path/to/the/file.jpg';
-
-  it('should return expected values at initialization', async () => {
-    (Storage.get as jest.Mock).mockResolvedValue(undefined);
-
-    const { result, waitForNextUpdate } = renderHook(() =>
-      useStorageURL(storageKey)
-    );
-
-    expect(result.current.error).toBeUndefined();
-    expect(result.current.isLoading).toBe(true);
-    expect(result.current.url).toBeUndefined();
-
-    // Force next render to prevent test warning
-    await waitForNextUpdate();
-  });
+  const errorConfig = {
+    fallbackURL: 'https://amplify.s3.amazonaws.com/path/to/the/fallback.jpg',
+    onStorageGetError: jest.fn(),
+  };
 
   it('should return a Storage URL', async () => {
     (Storage.get as jest.Mock).mockResolvedValue(storageUrl);
 
     const { result, waitForNextUpdate } = renderHook(() =>
-      useStorageURL(storageKey)
+      useStorageURL(storageKey, storageOptions)
     );
 
-    await waitForNextUpdate();
-
-    expect(result.current.error).toBeUndefined();
-    expect(result.current.isLoading).toBe(false);
-    expect(result.current.url).toBe(storageUrl);
-  });
-
-  it('should invoke Storage.get on fetch', async () => {
-    const mockStorageGet = jest.fn(() => Promise.resolve());
-
-    (Storage.get as jest.Mock).mockImplementation(mockStorageGet);
+    // should return undefined at initialization
+    expect(result.current).toBeUndefined();
 
-    const { waitForNextUpdate } = renderHook(() =>
-      useStorageURL(storageKey, storageOptions)
-    );
+    expect(Storage.get).toHaveBeenCalledWith(storageKey, storageOptions);
 
+    // Next update will happen when Storage.get resolves
     await waitForNextUpdate();
 
-    expect(mockStorageGet).toHaveBeenCalledWith(storageKey, storageOptions);
+    expect(result.current).toBe(storageUrl);
+    expect(Storage.get).toHaveBeenCalledTimes(1);
   });
 
-  it('should set an error when Storage.get fails', async () => {
+  it('should invoke onStorageGetError and return a fallbackURL when Storage.get fails', async () => {
     const customError = new Error('Something wrong happen');
 
     (Storage.get as jest.Mock).mockRejectedValue(customError);
 
     const { result, waitForNextUpdate } = renderHook(() =>
-      useStorageURL(storageKey)
+      useStorageURL(storageKey, storageOptions, errorConfig)
     );
 
+    expect(Storage.get).toHaveBeenCalledWith(storageKey, storageOptions);
+
+    // Next update will happen when Storage.get resolves
     await waitForNextUpdate();
 
-    expect(result.current.error).toBe(customError);
-    expect(result.current.isLoading).toBe(false);
-    expect(result.current.url).toBeUndefined();
+    expect(result.current).toBe(errorConfig.fallbackURL);
+    expect(errorConfig.onStorageGetError).toHaveBeenCalledTimes(1);
+    expect(errorConfig.onStorageGetError).toHaveBeenCalledWith(customError);
   });
 
-  it('should cancel Storage.get request on unmount', async () => {
-    const mockStorageCancel = jest.fn();
-
-    (Storage.get as jest.Mock).mockResolvedValue(undefined);
-    (Storage.cancel as jest.Mock).mockImplementation(mockStorageCancel);
+  it('should execute Storage.cancel before rerendering', async () => {
+    (Storage.get as jest.Mock).mockResolvedValue(storageUrl);
 
-    const { waitForNextUpdate, unmount } = renderHook(() =>
-      useStorageURL(storageKey)
+    const { waitForNextUpdate } = renderHook(() =>
+      useStorageURL(storageKey, storageOptions)
     );
 
-    // Start Storage fetch
-    await waitForNextUpdate();
+    expect(Storage.get).toHaveBeenCalledWith(storageKey, storageOptions);
 
-    // Unmount!
-    act(() => unmount());
+    // Next update will happen when Storage.get resolves
+    await waitForNextUpdate();
 
-    expect(mockStorageCancel).toHaveBeenCalled();
+    // Since a rerender has happened, Storage.cancel should be run once at this point as a useEffect cleanup function
+    expect(Storage.cancel).toHaveBeenCalled();
+    expect(Storage.cancel).toHaveBeenCalledTimes(1);
   });
 });

packages/react/src/hooks/useStorageURL.ts

@@ -1,11 +1,12 @@
 import * as React from 'react';
 
+import { isFunction } from '@aws-amplify/ui';
+import { useHasValueUpdated } from '@aws-amplify/ui-react-core';
 import { S3ProviderGetConfig, Storage } from '@aws-amplify/storage';
 
-export interface UseStorageURLResult {
-  url?: string;
-  error?: Error;
-  isLoading: boolean;
+interface UseStorageURLErrorConfig {
+  fallbackURL?: string;
+  onStorageGetError?: (error: Error) => void;
 }
 
 /**
@@ -14,31 +15,32 @@ export interface UseStorageURLResult {
  */
 export const useStorageURL = (
   key: string,
-  options?: S3ProviderGetConfig
-): UseStorageURLResult & { fetch: () => () => void } => {
-  const [result, setResult] = React.useState<UseStorageURLResult>({
-    isLoading: true,
-  });
-
-  // Used to prevent an infinite loop on useEffect, because `options`
-  // will have a different reference on every render
-  const serializedOptions = JSON.stringify(options);
-
-  const fetch = () => {
-    setResult({ isLoading: true });
-
-    const promise = Storage.get(key, options);
-
-    // Attempt to fetch storage object url
-    promise
-      .then((url) => setResult({ url, isLoading: false }))
-      .catch((error: Error) => setResult({ error, isLoading: false }));
+  options?: S3ProviderGetConfig,
+  errorConfig?: UseStorageURLErrorConfig
+): string | undefined => {
+  const [url, setURL] = React.useState<string>();
+  const hasKeyUpdated = useHasValueUpdated(key);
+
+  React.useEffect(() => {
+    if (!hasKeyUpdated) {
+      return;
+    }
+
+    const promise = Storage.get(key, options)
+      .then((url) => setURL(url))
+      .catch((error: Error) => {
+        const { fallbackURL, onStorageGetError } = errorConfig ?? {};
+        if (isFunction(onStorageGetError)) {
+          onStorageGetError(error);
+        }
+        if (fallbackURL) {
+          setURL(fallbackURL);
+        }
+      });
 
     // Cancel current promise on unmount
     return () => Storage.cancel(promise);
-  };
-
-  React.useEffect(fetch, [key, options, serializedOptions]);
+  }, [key, options, errorConfig, hasKeyUpdated]);
 
-  return { ...result, fetch };
+  return url;
 };

packages/react/src/primitives/shared/constants.ts

@@ -548,6 +548,10 @@ export const ComponentClassObject: ComponentClassNameItems = {
     components: ['StepperField'],
     description: 'Class applied to the StepperField input',
   },
+  StorageImage: {
+    className: ComponentClassName.StorageImage,
+    components: ['StorageImage'],
+  },
   StorageManager: {
     className: ComponentClassName.StorageManager,
     components: ['StorageManager'],
@@ -868,6 +872,7 @@ export const ComponentClassNames: ComponentClassNamesType = {
   StepperFieldButtonIncrease:
     ComponentClassObject.StepperFieldButtonIncrease.className,
   StepperFieldInput: ComponentClassObject.StepperFieldInput.className,
+  StorageImage: ComponentClassObject.StorageImage.className,
   StorageManager: ComponentClassObject.StorageManager.className,
   StorageManagerDropZone: ComponentClassObject.StorageManagerDropZone.className,
   StorageManagerDropZoneIcon: