Update per review

leo-generali · leo-generali · commit 615c56638249 · 2025-09-08T14:53:09.000-04:00
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx
@@ -319,7 +319,7 @@ describe('MockDataGeneratorModal', () => {
         .to.not.exist;
     });
 
-    // todo: assert that the generated script is displayed in the code block
+    // todo: assert that the generated script is displayed in the code block (CLOUDP-333860)
   });
 
   describe('when rendering the modal in a specific step', () => {
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useMemo } from 'react';
 import { connect } from 'react-redux';
 
 import {
@@ -25,14 +25,6 @@ import { default as SchemaConfirmationScreen } from './raw-schema-confirmation';
 import FakerSchemaEditor from './faker-schema-editor';
 import ScriptScreen from './script-screen';
 
-const STEP_TO_STEP_CONTENT: Record<MockDataGeneratorStep, React.JSX.Element> = {
-  [MockDataGeneratorStep.SCHEMA_CONFIRMATION]: <SchemaConfirmationScreen />,
-  [MockDataGeneratorStep.SCHEMA_EDITOR]: <FakerSchemaEditor />,
-  [MockDataGeneratorStep.DOCUMENT_COUNT]: <></>, // TODO: Implement as part of CLOUDP-333856
-  [MockDataGeneratorStep.PREVIEW_DATA]: <></>, // TODO: Implement as part of CLOUDP-333857
-  [MockDataGeneratorStep.GENERATE_DATA]: <ScriptScreen />,
-};
-
 const footerStyles = css`
   flex-direction: row;
   justify-content: space-between;
@@ -61,6 +53,21 @@ const MockDataGeneratorModal = ({
   onConfirmSchema,
   onPreviousStep,
 }: Props) => {
+  const modalBodyContent = useMemo(() => {
+    switch (currentStep) {
+      case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
+        return <SchemaConfirmationScreen />;
+      case MockDataGeneratorStep.SCHEMA_EDITOR:
+        return <FakerSchemaEditor />;
+      case MockDataGeneratorStep.DOCUMENT_COUNT:
+        return <></>; // TODO: CLOUDP-333856
+      case MockDataGeneratorStep.PREVIEW_DATA:
+        return <></>; // TODO: CLOUDP-333857
+      case MockDataGeneratorStep.GENERATE_DATA:
+        return <ScriptScreen />;
+    }
+  }, [currentStep]);
+
   const handleNextClick = () => {
     if (currentStep === MockDataGeneratorStep.GENERATE_DATA) {
       onClose();
@@ -85,7 +92,7 @@ const MockDataGeneratorModal = ({
       <ModalHeader title="Generate Mock Data" />
       <ModalBody>
         <div data-testid={`generate-mock-data-step-${currentStep}`}>
-          {STEP_TO_STEP_CONTENT[currentStep]}
+          {modalBodyContent}
         </div>
       </ModalBody>
       <ModalFooter className={footerStyles}>
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx
@@ -47,6 +47,10 @@ const resourceSectionStyles = css({
   borderRadius: spacing[400],
 });
 
+const instructionTextStyle = css({
+  marginTop: spacing[200],
+});
+
 const resourceSectionLightStyles = css({
   backgroundColor: palette.gray.light3,
 });
@@ -69,7 +73,7 @@ const ScriptScreen = () => {
         <Body as="h2" baseFontSize={16} weight="medium">
           Prerequisites
         </Body>
-        <Body className={css({ marginTop: spacing[200] })}>
+        <Body className={instructionTextStyle}>
           To run the generated script, you must:
         </Body>
         <ul className={listStyles}>
