[Cases] Improve attachments rendering (#139827)

x-pack/plugins/cases/public/components/user_actions/comment/comment.test.tsx

@@ -7,7 +7,7 @@
 
 import React from 'react';
 import { EuiCommentList } from '@elastic/eui';
-import { render, screen } from '@testing-library/react';
+import { render, screen, waitFor } from '@testing-library/react';
 
 import { Actions } from '../../../../common/api';
 import {
@@ -273,27 +273,78 @@ describe('createCommentUserActionBuilder', () => {
     });
 
     it('renders correctly a persistable state attachment', async () => {
+      const MockComponent = jest.fn((props) => {
+        return (
+          <div data-test-subj={`attachment_${props.persistableStateAttachmentState.test_foo}`} />
+        );
+      });
+
+      const SpyLazyFactory = jest.fn(() => {
+        return Promise.resolve().then(() => {
+          return {
+            default: React.memo(MockComponent),
+          };
+        });
+      });
+
       const persistableStateAttachmentTypeRegistry = new PersistableStateAttachmentTypeRegistry();
-      persistableStateAttachmentTypeRegistry.register(getPersistableStateAttachment());
+      persistableStateAttachmentTypeRegistry.register(
+        getPersistableStateAttachment({
+          children: React.lazy(SpyLazyFactory),
+        })
+      );
 
       const userAction = getPersistableStateUserAction();
+      const attachment01 = {
+        ...persistableStateAttachment,
+        persistableStateAttachmentState: { test_foo: '01' },
+      };
       const builder = createCommentUserActionBuilder({
         ...builderArgs,
         persistableStateAttachmentTypeRegistry,
         caseData: {
           ...builderArgs.caseData,
-          comments: [persistableStateAttachment],
+          comments: [attachment01],
         },
         userAction,
       });
 
-      const createdUserAction = builder.build();
-      const result = appMockRender.render(<EuiCommentList comments={createdUserAction} />);
+      const result = appMockRender.render(<EuiCommentList comments={builder.build()} />);
+
+      await waitFor(() => {
+        expect(result.getByTestId('attachment_01')).toBeInTheDocument();
+        expect(MockComponent).toHaveBeenCalledTimes(1);
+        expect(SpyLazyFactory).toHaveBeenCalledTimes(1);
+      });
 
       expect(result.getByTestId('comment-persistableState-.test')).toBeInTheDocument();
       expect(result.getByTestId('copy-link-persistable-state-comment-id')).toBeInTheDocument();
       expect(result.getByTestId('user-action-username-with-avatar')).toBeInTheDocument();
       expect(screen.getByText('added an embeddable')).toBeInTheDocument();
+
+      result.unmount();
+
+      const attachment02 = {
+        ...persistableStateAttachment,
+        persistableStateAttachmentState: { test_foo: '02' },
+      };
+      const updateBuilder = createCommentUserActionBuilder({
+        ...builderArgs,
+        persistableStateAttachmentTypeRegistry,
+        caseData: {
+          ...builderArgs.caseData,
+          comments: [attachment02],
+        },
+        userAction,
+      });
+
+      const result2 = appMockRender.render(<EuiCommentList comments={updateBuilder.build()} />);
+
+      await waitFor(() => {
+        expect(result2.getByTestId('attachment_02')).toBeInTheDocument();
+        expect(MockComponent).toHaveBeenCalledTimes(2);
+        expect(SpyLazyFactory).toHaveBeenCalledTimes(1);
+      });
     });
 
     it('renders correctly if the reference is not registered', async () => {

x-pack/plugins/cases/public/components/user_actions/comment/registered_attachments.tsx

@@ -12,6 +12,7 @@
  */
 
 import React, { Suspense } from 'react';
+import { memoize } from 'lodash';
 
 import { EuiCallOut, EuiCode, EuiLoadingSpinner } from '@elastic/eui';
 import { AttachmentType } from '../../../client/attachment_framework/types';
@@ -31,6 +32,29 @@ type BuilderArgs<C, R> = Pick<UserActionBuilderArgs, 'userAction' | 'caseData'>
   getAttachmentViewProps: () => object;
 };
 
+/**
+ * Provides a render function for attachment type
+ */
+const getAttachmentRenderer = memoize((attachmentType: AttachmentType<unknown>) => {
+  const attachmentViewObject = attachmentType.getAttachmentViewObject();
+
+  let AttachmentElement: React.ReactElement;
+
+  const renderCallback = (props: object) => {
+    if (!attachmentViewObject.children) return;
+
+    if (!AttachmentElement) {
+      AttachmentElement = React.createElement(attachmentViewObject.children, props);
+    } else {
+      AttachmentElement = React.cloneElement(AttachmentElement, props);
+    }
+
+    return <Suspense fallback={<EuiLoadingSpinner />}>{AttachmentElement}</Suspense>;
+  };
+
+  return renderCallback;
+});
+
 export const createRegisteredAttachmentUserActionBuilder = <
   C extends CommentResponse,
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -75,6 +99,7 @@ export const createRegisteredAttachmentUserActionBuilder = <
     }
 
     const attachmentType = registry.get(attachmentTypeId);
+    const renderer = getAttachmentRenderer(attachmentType);
 
     const attachmentViewObject = attachmentType.getAttachmentViewObject();
     const props = {
@@ -101,11 +126,7 @@ export const createRegisteredAttachmentUserActionBuilder = <
             {attachmentViewObject.actions}
           </>
         ),
-        children: attachmentViewObject.children ? (
-          <Suspense fallback={<EuiLoadingSpinner />}>
-            {React.createElement(attachmentViewObject.children, props)}
-          </Suspense>
-        ) : undefined,
+        children: renderer(props),
       },
     ];
   },