From 8e0240184f790ce4125f37a02b24dc27677b13ef Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Wed, 30 Jul 2025 13:12:28 -0400
Subject: [PATCH 1/9] WIP

---
 .../mock-data-generator-modal/constants.ts    |   1 +
 .../mock-generator-modal.tsx                  | 104 ++++++++++++++++++
 .../mock-data-generator-modal/types.ts        |   8 ++
 3 files changed, 113 insertions(+)
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/types.ts

diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
new file mode 100644
index 00000000000..0068109918f
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
@@ -0,0 +1 @@
+export const DEFAULT_OUTPUT_DOCS_COUNT = 1000;
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx
new file mode 100644
index 00000000000..3d42da426f9
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx
@@ -0,0 +1,104 @@
+import React, { useState } from 'react';
+
+import { css } from '@mongodb-js/compass-components';
+
+import {
+  Button,
+  Modal,
+  ModalHeader,
+  ModalBody,
+  ModalFooter,
+  ButtonVariant,
+} from '@mongodb-js/compass-components';
+import { MockDataGeneratorSteps } from './types';
+import { DEFAULT_OUTPUT_DOCS_COUNT } from './constants';
+
+const footerStyles = css`
+  flex-direction: row;
+  justify-content: space-between;
+`;
+
+const rightButtonsStyles = css`
+  display: flex;
+  gap: 8px;
+  flex-direction: row;
+`;
+
+interface Props {
+  isOpen: boolean;
+  setIsOpen: (isOpen: boolean) => void;
+}
+
+const MockDataGeneratorModal = ({ isOpen, setIsOpen }: Props) => {
+  const [currentStep, setCurrentStep] = useState<MockDataGeneratorSteps>(
+    MockDataGeneratorSteps.AI_DISCLAIMER
+  );
+  const [rawSchema, setRawSchema] = useState<string | null>(null);
+  const [fakerSchema, setFakerSchema] = useState<string | null>(null);
+  const [outputDocsCount, setOutputDocsCount] = useState<number>(
+    DEFAULT_OUTPUT_DOCS_COUNT
+  );
+  const [validationRules, setValidationRules] = useState<string | null>(null);
+  const [sampleDoc, setSampleDoc] = useState<string | null>(null);
+
+  const resetState = () => {
+    setCurrentStep(MockDataGeneratorSteps.AI_DISCLAIMER);
+    setRawSchema(null);
+    setFakerSchema(null);
+    setOutputDocsCount(DEFAULT_OUTPUT_DOCS_COUNT);
+    setValidationRules(null);
+    setSampleDoc(null);
+  };
+
+  const onNext = () => {
+    if (currentStep < MockDataGeneratorSteps.GENERATE_DATA) {
+      setCurrentStep(currentStep + 1);
+    } else {
+      // Final step, close the modal
+      setIsOpen(false);
+      resetState();
+    }
+  };
+
+  const onBack = () => {
+    if (currentStep > MockDataGeneratorSteps.AI_DISCLAIMER) {
+      setCurrentStep(currentStep - 1);
+    }
+  };
+
+  const onCancel = () => {
+    setIsOpen(false);
+  };
+
+  return (
+    <Modal
+      open={isOpen}
+      setOpen={() => setIsOpen(false)}
+      data-testid="generate-mock-data-modal"
+    >
+      <ModalHeader title="Generate Mock Data" />
+      <ModalBody>
+        {currentStep === MockDataGeneratorSteps.AI_DISCLAIMER && <div></div>}
+        {currentStep === MockDataGeneratorSteps.SCHEMA_CONFIRMATION && (
+          <div></div>
+        )}
+        {currentStep === MockDataGeneratorSteps.SCHEMA_EDITOR && <div></div>}
+        {currentStep === MockDataGeneratorSteps.DOCUMENT_COUNT && <div></div>}
+        {currentStep === MockDataGeneratorSteps.PREVIEW_DATA && <div></div>}
+        {currentStep === MockDataGeneratorSteps.GENERATE_DATA && <div></div>}
+      </ModalBody>
+
+      <ModalFooter className={footerStyles}>
+        <Button onClick={onBack}>Back</Button>
+        <div className={rightButtonsStyles}>
+          <Button onClick={onCancel}>Cancel</Button>
+          <Button variant={ButtonVariant.Primary} onClick={onNext}>
+            Next
+          </Button>
+        </div>
+      </ModalFooter>
+    </Modal>
+  );
+};
+
+export default MockDataGeneratorModal;
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
new file mode 100644
index 00000000000..c1ca93a830c
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
@@ -0,0 +1,8 @@
+export enum MockDataGeneratorSteps {
+  AI_DISCLAIMER = 1,
+  SCHEMA_CONFIRMATION = 2,
+  SCHEMA_EDITOR = 3,
+  DOCUMENT_COUNT = 4,
+  PREVIEW_DATA = 5,
+  GENERATE_DATA = 6,
+}

From 848d59ce00b630b696533edd11e9e770d9bb2837 Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Thu, 31 Jul 2025 13:10:57 -0400
Subject: [PATCH 2/9] add unit tests

---
 .../mock-data-generator-modal/constants.ts    |  1 -
 .../mock-data-generator-modal.spec.tsx        | 57 +++++++++++++++++++
 ...odal.tsx => mock-data-generator-modal.tsx} | 28 +--------
 3 files changed, 59 insertions(+), 27 deletions(-)
 delete mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx
 rename packages/compass-collection/src/components/mock-data-generator-modal/{mock-generator-modal.tsx => mock-data-generator-modal.tsx} (58%)

diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
deleted file mode 100644
index 0068109918f..00000000000
--- a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
+++ /dev/null
@@ -1 +0,0 @@
-export const DEFAULT_OUTPUT_DOCS_COUNT = 1000;
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
new file mode 100644
index 00000000000..b2b24b6a6ec
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx
@@ -0,0 +1,57 @@
+import { expect } from 'chai';
+import React from 'react';
+import { render, screen } from '@mongodb-js/testing-library-compass';
+import Sinon from 'sinon';
+import MockDataGeneratorModal from './mock-data-generator-modal';
+import { MockDataGeneratorSteps } from './types.ts';
+
+describe('MockDataGeneratorModal', function () {
+  const sandbox = Sinon.createSandbox();
+  let setIsOpen: Sinon.SinonSpy;
+
+  beforeEach(function () {
+    setIsOpen = sandbox.spy();
+  });
+
+  afterEach(function () {
+    sandbox.restore();
+  });
+
+  const renderModal = () => {
+    return render(
+      <MockDataGeneratorModal isOpen={true} setIsOpen={setIsOpen} />
+    );
+  };
+
+  it('renders the modal when isOpen is true', function () {
+    renderModal();
+
+    expect(screen.getByTestId('generate-mock-data-modal')).to.exist;
+  });
+
+  it('calls setIsOpen(false) when the modal is closed', function () {
+    renderModal();
+
+    screen.getByLabelText('Close modal').click();
+
+    expect(setIsOpen.calledOnceWith(false)).to.be.true;
+  });
+
+  it('calls setIsOpen(false) when the cancel button is clicked', function () {
+    renderModal();
+
+    screen.getByText('Cancel').click();
+
+    expect(setIsOpen.calledOnceWith(false)).to.be.true;
+  });
+
+  it('renders the first step by default', function () {
+    renderModal();
+
+    expect(
+      screen.getByTestId(
+        `generate-mock-data-step-${MockDataGeneratorSteps.AI_DISCLAIMER}`
+      )
+    ).to.exist;
+  });
+});
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx
similarity index 58%
rename from packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx
rename to packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx
index 3d42da426f9..f9a1b41daa9 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/mock-generator-modal.tsx
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx
@@ -1,17 +1,14 @@
 import React, { useState } from 'react';
 
-import { css } from '@mongodb-js/compass-components';
+import { css, ModalBody, ModalHeader } from '@mongodb-js/compass-components';
 
 import {
   Button,
   Modal,
-  ModalHeader,
-  ModalBody,
   ModalFooter,
   ButtonVariant,
 } from '@mongodb-js/compass-components';
 import { MockDataGeneratorSteps } from './types';
-import { DEFAULT_OUTPUT_DOCS_COUNT } from './constants';
 
 const footerStyles = css`
   flex-direction: row;
@@ -33,28 +30,15 @@ const MockDataGeneratorModal = ({ isOpen, setIsOpen }: Props) => {
   const [currentStep, setCurrentStep] = useState<MockDataGeneratorSteps>(
     MockDataGeneratorSteps.AI_DISCLAIMER
   );
-  const [rawSchema, setRawSchema] = useState<string | null>(null);
-  const [fakerSchema, setFakerSchema] = useState<string | null>(null);
-  const [outputDocsCount, setOutputDocsCount] = useState<number>(
-    DEFAULT_OUTPUT_DOCS_COUNT
-  );
-  const [validationRules, setValidationRules] = useState<string | null>(null);
-  const [sampleDoc, setSampleDoc] = useState<string | null>(null);
 
   const resetState = () => {
     setCurrentStep(MockDataGeneratorSteps.AI_DISCLAIMER);
-    setRawSchema(null);
-    setFakerSchema(null);
-    setOutputDocsCount(DEFAULT_OUTPUT_DOCS_COUNT);
-    setValidationRules(null);
-    setSampleDoc(null);
   };
 
   const onNext = () => {
     if (currentStep < MockDataGeneratorSteps.GENERATE_DATA) {
       setCurrentStep(currentStep + 1);
     } else {
-      // Final step, close the modal
       setIsOpen(false);
       resetState();
     }
@@ -78,16 +62,8 @@ const MockDataGeneratorModal = ({ isOpen, setIsOpen }: Props) => {
     >
       <ModalHeader title="Generate Mock Data" />
       <ModalBody>
-        {currentStep === MockDataGeneratorSteps.AI_DISCLAIMER && <div></div>}
-        {currentStep === MockDataGeneratorSteps.SCHEMA_CONFIRMATION && (
-          <div></div>
-        )}
-        {currentStep === MockDataGeneratorSteps.SCHEMA_EDITOR && <div></div>}
-        {currentStep === MockDataGeneratorSteps.DOCUMENT_COUNT && <div></div>}
-        {currentStep === MockDataGeneratorSteps.PREVIEW_DATA && <div></div>}
-        {currentStep === MockDataGeneratorSteps.GENERATE_DATA && <div></div>}
+        <div data-testid={`generate-mock-data-step-${currentStep}`} />
       </ModalBody>
-
       <ModalFooter className={footerStyles}>
         <Button onClick={onBack}>Back</Button>
         <div className={rightButtonsStyles}>

From fd608df9ae2b0dded86e2473530c069fa5a3d77c Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Thu, 31 Jul 2025 14:36:40 -0400
Subject: [PATCH 3/9] address comments

---
 .../mock-data-generator-modal.spec.tsx        | 87 ++++++++++++++++---
 .../mock-data-generator-modal.tsx             | 43 ++++++---
 .../mock-data-generator-modal/types.ts        |  2 +-
 .../mock-data-generator-modal/utils.ts        | 19 ++++
 4 files changed, 122 insertions(+), 29 deletions(-)
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/utils.ts

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
index b2b24b6a6ec..cb23e10d690 100644
--- 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
@@ -3,33 +3,64 @@ import React from 'react';
 import { render, screen } from '@mongodb-js/testing-library-compass';
 import Sinon from 'sinon';
 import MockDataGeneratorModal from './mock-data-generator-modal';
-import { MockDataGeneratorSteps } from './types.ts';
+import { MockDataGeneratorStep } from './types';
 
-describe('MockDataGeneratorModal', function () {
+describe('MockDataGeneratorModal', () => {
   const sandbox = Sinon.createSandbox();
   let setIsOpen: Sinon.SinonSpy;
 
-  beforeEach(function () {
+  beforeEach(() => {
     setIsOpen = sandbox.spy();
   });
 
-  afterEach(function () {
+  afterEach(() => {
     sandbox.restore();
   });
 
-  const renderModal = () => {
-    return render(
-      <MockDataGeneratorModal isOpen={true} setIsOpen={setIsOpen} />
-    );
-  };
+  function renderModal({
+    isOpen = true,
+    currentStep = MockDataGeneratorStep.AI_DISCLAIMER,
+  } = {}) {
+    function MockDataGeneratorModalWrapper() {
+      const [currentStepStateMock, setCurrentStepStateMock] =
+        React.useState<MockDataGeneratorStep>(currentStep);
+      return (
+        <MockDataGeneratorModal
+          isOpen={isOpen}
+          setIsOpen={setIsOpen}
+          currentStep={currentStepStateMock}
+          setCurrentStep={(step) => {
+            setCurrentStepStateMock(step);
+          }}
+        />
+      );
+    }
+    return render(<MockDataGeneratorModalWrapper />);
+  }
 
-  it('renders the modal when isOpen is true', function () {
+  it('renders the modal when isOpen is true', () => {
     renderModal();
 
     expect(screen.getByTestId('generate-mock-data-modal')).to.exist;
   });
 
-  it('calls setIsOpen(false) when the modal is closed', function () {
+  it('does not render the modal when isOpen is false', () => {
+    renderModal({ isOpen: false });
+
+    expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist;
+  });
+
+  it('renders the correct step when currentStep is set', () => {
+    renderModal({ currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION });
+
+    expect(
+      screen.getByTestId(
+        `generate-mock-data-step-${MockDataGeneratorStep.SCHEMA_CONFIRMATION}`
+      )
+    ).to.exist;
+  });
+
+  it('calls setIsOpen(false) when the modal is closed', () => {
     renderModal();
 
     screen.getByLabelText('Close modal').click();
@@ -37,7 +68,7 @@ describe('MockDataGeneratorModal', function () {
     expect(setIsOpen.calledOnceWith(false)).to.be.true;
   });
 
-  it('calls setIsOpen(false) when the cancel button is clicked', function () {
+  it('calls setIsOpen(false) when the cancel button is clicked', () => {
     renderModal();
 
     screen.getByText('Cancel').click();
@@ -45,12 +76,40 @@ describe('MockDataGeneratorModal', function () {
     expect(setIsOpen.calledOnceWith(false)).to.be.true;
   });
 
-  it('renders the first step by default', function () {
+  it('disables the Back button on the first step', () => {
+    renderModal();
+
+    expect(
+      screen.getByRole('button', { name: 'Back' }).getAttribute('aria-disabled')
+    ).to.equal('true');
+  });
+
+  it('renders the next step button with the correct label on each step', () => {
+    renderModal();
+
+    expect(screen.getByTestId('next-step-button')).to.have.text(
+      'Use Natural Language'
+    );
+    screen.getByTestId('next-step-button').click();
+    expect(screen.getByTestId('next-step-button')).to.have.text('Confirm');
+    screen.getByTestId('next-step-button').click();
+    expect(screen.getByTestId('next-step-button')).to.have.text('Next');
+    screen.getByTestId('next-step-button').click();
+    expect(screen.getByTestId('next-step-button')).to.have.text('Next');
+    screen.getByTestId('next-step-button').click();
+    expect(screen.getByTestId('next-step-button')).to.have.text(
+      'Generate Script'
+    );
+    screen.getByTestId('next-step-button').click();
+    expect(screen.getByTestId('next-step-button')).to.have.text('Done');
+  });
+
+  it('renders the first step by default', () => {
     renderModal();
 
     expect(
       screen.getByTestId(
-        `generate-mock-data-step-${MockDataGeneratorSteps.AI_DISCLAIMER}`
+        `generate-mock-data-step-${MockDataGeneratorStep.AI_DISCLAIMER}`
       )
     ).to.exist;
   });
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
index f9a1b41daa9..8fe9e1b888b 100644
--- 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, { useState } from 'react';
+import React from 'react';
 
 import { css, ModalBody, ModalHeader } from '@mongodb-js/compass-components';
 
@@ -8,7 +8,8 @@ import {
   ModalFooter,
   ButtonVariant,
 } from '@mongodb-js/compass-components';
-import { MockDataGeneratorSteps } from './types';
+import { MockDataGeneratorStep } from './types';
+import { getNextStepButtonLabel } from './utils';
 
 const footerStyles = css`
   flex-direction: row;
@@ -24,19 +25,22 @@ const rightButtonsStyles = css`
 interface Props {
   isOpen: boolean;
   setIsOpen: (isOpen: boolean) => void;
+  currentStep: MockDataGeneratorStep;
+  setCurrentStep: (step: MockDataGeneratorStep) => void;
 }
 
-const MockDataGeneratorModal = ({ isOpen, setIsOpen }: Props) => {
-  const [currentStep, setCurrentStep] = useState<MockDataGeneratorSteps>(
-    MockDataGeneratorSteps.AI_DISCLAIMER
-  );
-
+const MockDataGeneratorModal = ({
+  isOpen,
+  setIsOpen,
+  currentStep,
+  setCurrentStep,
+}: Props) => {
   const resetState = () => {
-    setCurrentStep(MockDataGeneratorSteps.AI_DISCLAIMER);
+    setCurrentStep(MockDataGeneratorStep.AI_DISCLAIMER);
   };
 
   const onNext = () => {
-    if (currentStep < MockDataGeneratorSteps.GENERATE_DATA) {
+    if (currentStep < MockDataGeneratorStep.GENERATE_DATA) {
       setCurrentStep(currentStep + 1);
     } else {
       setIsOpen(false);
@@ -45,31 +49,42 @@ const MockDataGeneratorModal = ({ isOpen, setIsOpen }: Props) => {
   };
 
   const onBack = () => {
-    if (currentStep > MockDataGeneratorSteps.AI_DISCLAIMER) {
+    if (currentStep > MockDataGeneratorStep.AI_DISCLAIMER) {
       setCurrentStep(currentStep - 1);
     }
   };
 
   const onCancel = () => {
     setIsOpen(false);
+    resetState();
   };
 
   return (
     <Modal
       open={isOpen}
-      setOpen={() => setIsOpen(false)}
+      setOpen={(open) => setIsOpen(open)}
       data-testid="generate-mock-data-modal"
     >
       <ModalHeader title="Generate Mock Data" />
       <ModalBody>
+        {/* TODO: Render actual step content here based on currentStep. (CLOUDP-333851) */}
         <div data-testid={`generate-mock-data-step-${currentStep}`} />
       </ModalBody>
       <ModalFooter className={footerStyles}>
-        <Button onClick={onBack}>Back</Button>
+        <Button
+          onClick={onBack}
+          disabled={currentStep === MockDataGeneratorStep.AI_DISCLAIMER}
+        >
+          Back
+        </Button>
         <div className={rightButtonsStyles}>
           <Button onClick={onCancel}>Cancel</Button>
-          <Button variant={ButtonVariant.Primary} onClick={onNext}>
-            Next
+          <Button
+            variant={ButtonVariant.Primary}
+            onClick={onNext}
+            data-testid="next-step-button"
+          >
+            {getNextStepButtonLabel(currentStep)}
           </Button>
         </div>
       </ModalFooter>
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
index c1ca93a830c..be15aaf2c65 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
@@ -1,4 +1,4 @@
-export enum MockDataGeneratorSteps {
+export enum MockDataGeneratorStep {
   AI_DISCLAIMER = 1,
   SCHEMA_CONFIRMATION = 2,
   SCHEMA_EDITOR = 3,
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts b/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
new file mode 100644
index 00000000000..bae7fcb41c7
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
@@ -0,0 +1,19 @@
+import { MockDataGeneratorStep } from './types';
+
+export const getNextStepButtonLabel = (step: MockDataGeneratorStep): string => {
+  switch (step) {
+    case MockDataGeneratorStep.AI_DISCLAIMER:
+      return 'Use Natural Language';
+    case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
+      return 'Confirm';
+    case MockDataGeneratorStep.SCHEMA_EDITOR:
+    case MockDataGeneratorStep.DOCUMENT_COUNT:
+      return 'Next';
+    case MockDataGeneratorStep.PREVIEW_DATA:
+      return 'Generate Script';
+    case MockDataGeneratorStep.GENERATE_DATA:
+      return 'Done';
+    default:
+      return 'Next';
+  }
+};

From 300cc0da5ef6b83319d67bcee5bf8d9a881c14d8 Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Thu, 31 Jul 2025 14:38:27 -0400
Subject: [PATCH 4/9] remove default step test

---
 .../mock-data-generator-modal.spec.tsx                 | 10 ----------
 1 file changed, 10 deletions(-)

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
index cb23e10d690..fc982d31b84 100644
--- 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
@@ -103,14 +103,4 @@ describe('MockDataGeneratorModal', () => {
     screen.getByTestId('next-step-button').click();
     expect(screen.getByTestId('next-step-button')).to.have.text('Done');
   });
-
-  it('renders the first step by default', () => {
-    renderModal();
-
-    expect(
-      screen.getByTestId(
-        `generate-mock-data-step-${MockDataGeneratorStep.AI_DISCLAIMER}`
-      )
-    ).to.exist;
-  });
 });

From 1804f3487655962a42dd886fd213b58201f609ad Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Thu, 31 Jul 2025 14:46:30 -0400
Subject: [PATCH 5/9] use constant instead of util function for getting button
 label

---
 .../mock-data-generator-modal/constants.ts    | 10 ++++++++++
 .../mock-data-generator-modal.spec.tsx        | 20 +++++--------------
 .../mock-data-generator-modal.tsx             |  4 ++--
 .../mock-data-generator-modal/utils.ts        | 19 ------------------
 4 files changed, 17 insertions(+), 36 deletions(-)
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
 delete mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/utils.ts

diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
new file mode 100644
index 00000000000..5057ee5946a
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts
@@ -0,0 +1,10 @@
+import { MockDataGeneratorStep } from './types';
+
+export const StepButtonLabelMap = {
+  [MockDataGeneratorStep.AI_DISCLAIMER]: 'Use Natural Language',
+  [MockDataGeneratorStep.SCHEMA_CONFIRMATION]: 'Confirm',
+  [MockDataGeneratorStep.SCHEMA_EDITOR]: 'Next',
+  [MockDataGeneratorStep.DOCUMENT_COUNT]: 'Next',
+  [MockDataGeneratorStep.PREVIEW_DATA]: 'Generate Script',
+  [MockDataGeneratorStep.GENERATE_DATA]: 'Done',
+} as const;
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
index fc982d31b84..9a7e249c1f2 100644
--- 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
@@ -4,6 +4,7 @@ import { render, screen } from '@mongodb-js/testing-library-compass';
 import Sinon from 'sinon';
 import MockDataGeneratorModal from './mock-data-generator-modal';
 import { MockDataGeneratorStep } from './types';
+import { StepButtonLabelMap } from './constants';
 
 describe('MockDataGeneratorModal', () => {
   const sandbox = Sinon.createSandbox();
@@ -87,20 +88,9 @@ describe('MockDataGeneratorModal', () => {
   it('renders the next step button with the correct label on each step', () => {
     renderModal();
 
-    expect(screen.getByTestId('next-step-button')).to.have.text(
-      'Use Natural Language'
-    );
-    screen.getByTestId('next-step-button').click();
-    expect(screen.getByTestId('next-step-button')).to.have.text('Confirm');
-    screen.getByTestId('next-step-button').click();
-    expect(screen.getByTestId('next-step-button')).to.have.text('Next');
-    screen.getByTestId('next-step-button').click();
-    expect(screen.getByTestId('next-step-button')).to.have.text('Next');
-    screen.getByTestId('next-step-button').click();
-    expect(screen.getByTestId('next-step-button')).to.have.text(
-      'Generate Script'
-    );
-    screen.getByTestId('next-step-button').click();
-    expect(screen.getByTestId('next-step-button')).to.have.text('Done');
+    Object.values(StepButtonLabelMap).forEach((label) => {
+      expect(screen.getByTestId('next-step-button')).to.have.text(label);
+      screen.getByTestId('next-step-button').click();
+    });
   });
 });
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
index 8fe9e1b888b..80f28f3c23d 100644
--- 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
@@ -9,7 +9,7 @@ import {
   ButtonVariant,
 } from '@mongodb-js/compass-components';
 import { MockDataGeneratorStep } from './types';
-import { getNextStepButtonLabel } from './utils';
+import { StepButtonLabelMap } from './constants';
 
 const footerStyles = css`
   flex-direction: row;
@@ -84,7 +84,7 @@ const MockDataGeneratorModal = ({
             onClick={onNext}
             data-testid="next-step-button"
           >
-            {getNextStepButtonLabel(currentStep)}
+            {StepButtonLabelMap[currentStep]}
           </Button>
         </div>
       </ModalFooter>
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts b/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
deleted file mode 100644
index bae7fcb41c7..00000000000
--- a/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
+++ /dev/null
@@ -1,19 +0,0 @@
-import { MockDataGeneratorStep } from './types';
-
-export const getNextStepButtonLabel = (step: MockDataGeneratorStep): string => {
-  switch (step) {
-    case MockDataGeneratorStep.AI_DISCLAIMER:
-      return 'Use Natural Language';
-    case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
-      return 'Confirm';
-    case MockDataGeneratorStep.SCHEMA_EDITOR:
-    case MockDataGeneratorStep.DOCUMENT_COUNT:
-      return 'Next';
-    case MockDataGeneratorStep.PREVIEW_DATA:
-      return 'Generate Script';
-    case MockDataGeneratorStep.GENERATE_DATA:
-      return 'Done';
-    default:
-      return 'Next';
-  }
-};

From b99c7eb29235d1d1c4096319196d8db27fc531e6 Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Thu, 31 Jul 2025 18:16:31 -0400
Subject: [PATCH 6/9] address comment: refactor unit test

---
 .../mock-data-generator-modal.spec.tsx        | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

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
index 9a7e249c1f2..296e6784883 100644
--- 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
@@ -51,16 +51,6 @@ describe('MockDataGeneratorModal', () => {
     expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist;
   });
 
-  it('renders the correct step when currentStep is set', () => {
-    renderModal({ currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION });
-
-    expect(
-      screen.getByTestId(
-        `generate-mock-data-step-${MockDataGeneratorStep.SCHEMA_CONFIRMATION}`
-      )
-    ).to.exist;
-  });
-
   it('calls setIsOpen(false) when the modal is closed', () => {
     renderModal();
 
@@ -85,12 +75,18 @@ describe('MockDataGeneratorModal', () => {
     ).to.equal('true');
   });
 
-  it('renders the next step button with the correct label on each step', () => {
-    renderModal();
-
-    Object.values(StepButtonLabelMap).forEach((label) => {
-      expect(screen.getByTestId('next-step-button')).to.have.text(label);
-      screen.getByTestId('next-step-button').click();
+  describe('when rendering the modal in a specific step', () => {
+    const steps = Object.values(MockDataGeneratorStep).filter(
+      (step) => typeof step === 'number'
+    ) as MockDataGeneratorStep[];
+
+    steps.forEach((currentStep) => {
+      it(`renders the button with the correct label when the user is in step "${currentStep}"`, () => {
+        renderModal({ currentStep });
+        expect(screen.getByTestId('next-step-button')).to.have.text(
+          StepButtonLabelMap[currentStep]
+        );
+      });
     });
   });
 });

From e36b798633a235efe4d6cca70f41b5b917c6098a Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Mon, 4 Aug 2025 11:52:56 -0400
Subject: [PATCH 7/9] address comments

---
 .../mock-data-generator-modal.spec.tsx        | 26 ++++++++---------
 .../mock-data-generator-modal.tsx             | 29 +++++++++++--------
 2 files changed, 30 insertions(+), 25 deletions(-)

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
index 296e6784883..88d94c6e778 100644
--- 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
@@ -8,10 +8,10 @@ import { StepButtonLabelMap } from './constants';
 
 describe('MockDataGeneratorModal', () => {
   const sandbox = Sinon.createSandbox();
-  let setIsOpen: Sinon.SinonSpy;
+  let onOpenChange: Sinon.SinonSpy;
 
   beforeEach(() => {
-    setIsOpen = sandbox.spy();
+    onOpenChange = sandbox.spy();
   });
 
   afterEach(() => {
@@ -23,15 +23,15 @@ describe('MockDataGeneratorModal', () => {
     currentStep = MockDataGeneratorStep.AI_DISCLAIMER,
   } = {}) {
     function MockDataGeneratorModalWrapper() {
-      const [currentStepStateMock, setCurrentStepStateMock] =
+      const [currentStepStateMock, onCurrentStepChangeStateMock] =
         React.useState<MockDataGeneratorStep>(currentStep);
       return (
         <MockDataGeneratorModal
           isOpen={isOpen}
-          setIsOpen={setIsOpen}
+          onOpenChange={onOpenChange}
           currentStep={currentStepStateMock}
-          setCurrentStep={(step) => {
-            setCurrentStepStateMock(step);
+          onCurrentStepChange={(step) => {
+            onCurrentStepChangeStateMock(step);
           }}
         />
       );
@@ -51,20 +51,20 @@ describe('MockDataGeneratorModal', () => {
     expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist;
   });
 
-  it('calls setIsOpen(false) when the modal is closed', () => {
+  it('calls onOpenChange(false) when the modal is closed', () => {
     renderModal();
 
     screen.getByLabelText('Close modal').click();
 
-    expect(setIsOpen.calledOnceWith(false)).to.be.true;
+    expect(onOpenChange.calledOnceWith(false)).to.be.true;
   });
 
-  it('calls setIsOpen(false) when the cancel button is clicked', () => {
+  it('calls onOpenChange(false) when the cancel button is clicked', () => {
     renderModal();
 
     screen.getByText('Cancel').click();
 
-    expect(setIsOpen.calledOnceWith(false)).to.be.true;
+    expect(onOpenChange.calledOnceWith(false)).to.be.true;
   });
 
   it('disables the Back button on the first step', () => {
@@ -76,9 +76,9 @@ describe('MockDataGeneratorModal', () => {
   });
 
   describe('when rendering the modal in a specific step', () => {
-    const steps = Object.values(MockDataGeneratorStep).filter(
-      (step) => typeof step === 'number'
-    ) as MockDataGeneratorStep[];
+    const steps = Object.keys(
+      StepButtonLabelMap
+    ) as unknown as MockDataGeneratorStep[];
 
     steps.forEach((currentStep) => {
       it(`renders the button with the correct label when the user is in step "${currentStep}"`, () => {
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
index 80f28f3c23d..728f045cfd9 100644
--- 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,6 +1,11 @@
 import React from 'react';
 
-import { css, ModalBody, ModalHeader } from '@mongodb-js/compass-components';
+import {
+  css,
+  ModalBody,
+  ModalHeader,
+  spacing,
+} from '@mongodb-js/compass-components';
 
 import {
   Button,
@@ -18,51 +23,51 @@ const footerStyles = css`
 
 const rightButtonsStyles = css`
   display: flex;
-  gap: 8px;
+  gap: ${spacing[200]}px;
   flex-direction: row;
 `;
 
 interface Props {
   isOpen: boolean;
-  setIsOpen: (isOpen: boolean) => void;
+  onOpenChange: (isOpen: boolean) => void;
   currentStep: MockDataGeneratorStep;
-  setCurrentStep: (step: MockDataGeneratorStep) => void;
+  onCurrentStepChange: (step: MockDataGeneratorStep) => void;
 }
 
 const MockDataGeneratorModal = ({
   isOpen,
-  setIsOpen,
+  onOpenChange,
   currentStep,
-  setCurrentStep,
+  onCurrentStepChange,
 }: Props) => {
   const resetState = () => {
-    setCurrentStep(MockDataGeneratorStep.AI_DISCLAIMER);
+    onCurrentStepChange(MockDataGeneratorStep.AI_DISCLAIMER);
   };
 
   const onNext = () => {
     if (currentStep < MockDataGeneratorStep.GENERATE_DATA) {
-      setCurrentStep(currentStep + 1);
+      onCurrentStepChange(currentStep + 1);
     } else {
-      setIsOpen(false);
+      onOpenChange(false);
       resetState();
     }
   };
 
   const onBack = () => {
     if (currentStep > MockDataGeneratorStep.AI_DISCLAIMER) {
-      setCurrentStep(currentStep - 1);
+      onCurrentStepChange(currentStep - 1);
     }
   };
 
   const onCancel = () => {
-    setIsOpen(false);
+    onOpenChange(false);
     resetState();
   };
 
   return (
     <Modal
       open={isOpen}
-      setOpen={(open) => setIsOpen(open)}
+      setOpen={(open) => onOpenChange(open)}
       data-testid="generate-mock-data-modal"
     >
       <ModalHeader title="Generate Mock Data" />

From e4ad4ca00a5acc18bb24261322ee7ecfe9d3b8d5 Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Mon, 4 Aug 2025 13:24:38 -0400
Subject: [PATCH 8/9] remove resetState

---
 .../mock-data-generator-modal/mock-data-generator-modal.tsx | 6 ------
 1 file changed, 6 deletions(-)

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
index 728f045cfd9..67a1df675c5 100644
--- 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
@@ -40,16 +40,11 @@ const MockDataGeneratorModal = ({
   currentStep,
   onCurrentStepChange,
 }: Props) => {
-  const resetState = () => {
-    onCurrentStepChange(MockDataGeneratorStep.AI_DISCLAIMER);
-  };
-
   const onNext = () => {
     if (currentStep < MockDataGeneratorStep.GENERATE_DATA) {
       onCurrentStepChange(currentStep + 1);
     } else {
       onOpenChange(false);
-      resetState();
     }
   };
 
@@ -61,7 +56,6 @@ const MockDataGeneratorModal = ({
 
   const onCancel = () => {
     onOpenChange(false);
-    resetState();
   };
 
   return (

From e08a786bc88936e37b64e17309fc278dbfeb47f1 Mon Sep 17 00:00:00 2001
From: Nataly Carbonell <nataly.carbonell@mongodb.com>
Date: Tue, 5 Aug 2025 09:46:41 -0400
Subject: [PATCH 9/9] refactor modal step enum and logic

---
 .../mock-data-generator-modal.tsx             | 13 ++-
 .../mock-data-generator-modal/types.ts        | 12 +--
 .../mock-data-generator-modal/utils.spec.ts   | 81 +++++++++++++++++++
 .../mock-data-generator-modal/utils.ts        | 41 ++++++++++
 4 files changed, 133 insertions(+), 14 deletions(-)
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/utils.spec.ts
 create mode 100644 packages/compass-collection/src/components/mock-data-generator-modal/utils.ts

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
index 67a1df675c5..b7723530d61 100644
--- 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
@@ -15,6 +15,7 @@ import {
 } from '@mongodb-js/compass-components';
 import { MockDataGeneratorStep } from './types';
 import { StepButtonLabelMap } from './constants';
+import { getNextStep, getPreviousStep } from './utils';
 
 const footerStyles = css`
   flex-direction: row;
@@ -41,17 +42,13 @@ const MockDataGeneratorModal = ({
   onCurrentStepChange,
 }: Props) => {
   const onNext = () => {
-    if (currentStep < MockDataGeneratorStep.GENERATE_DATA) {
-      onCurrentStepChange(currentStep + 1);
-    } else {
-      onOpenChange(false);
-    }
+    const nextStep = getNextStep(currentStep);
+    onCurrentStepChange(nextStep);
   };
 
   const onBack = () => {
-    if (currentStep > MockDataGeneratorStep.AI_DISCLAIMER) {
-      onCurrentStepChange(currentStep - 1);
-    }
+    const previousStep = getPreviousStep(currentStep);
+    onCurrentStepChange(previousStep);
   };
 
   const onCancel = () => {
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
index be15aaf2c65..6394dc25e2d 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts
@@ -1,8 +1,8 @@
 export enum MockDataGeneratorStep {
-  AI_DISCLAIMER = 1,
-  SCHEMA_CONFIRMATION = 2,
-  SCHEMA_EDITOR = 3,
-  DOCUMENT_COUNT = 4,
-  PREVIEW_DATA = 5,
-  GENERATE_DATA = 6,
+  AI_DISCLAIMER = 'AI_DISCLAIMER',
+  SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION',
+  SCHEMA_EDITOR = 'SCHEMA_EDITOR',
+  DOCUMENT_COUNT = 'DOCUMENT_COUNT',
+  PREVIEW_DATA = 'PREVIEW_DATA',
+  GENERATE_DATA = 'GENERATE_DATA',
 }
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/utils.spec.ts b/packages/compass-collection/src/components/mock-data-generator-modal/utils.spec.ts
new file mode 100644
index 00000000000..239d3ed2d1e
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/utils.spec.ts
@@ -0,0 +1,81 @@
+import { expect } from 'chai';
+import { getNextStep, getPreviousStep } from './utils';
+import { MockDataGeneratorStep } from './types';
+
+describe('Mock Data Generator Modal Utils', () => {
+  describe('getNextStep', () => {
+    it('should go from AI_DISCLAIMER to SCHEMA_CONFIRMATION', () => {
+      expect(getNextStep(MockDataGeneratorStep.AI_DISCLAIMER)).to.equal(
+        MockDataGeneratorStep.SCHEMA_CONFIRMATION
+      );
+    });
+
+    it('should go from SCHEMA_CONFIRMATION to SCHEMA_EDITOR', () => {
+      expect(getNextStep(MockDataGeneratorStep.SCHEMA_CONFIRMATION)).to.equal(
+        MockDataGeneratorStep.SCHEMA_EDITOR
+      );
+    });
+
+    it('should go from SCHEMA_EDITOR to DOCUMENT_COUNT', () => {
+      expect(getNextStep(MockDataGeneratorStep.SCHEMA_EDITOR)).to.equal(
+        MockDataGeneratorStep.DOCUMENT_COUNT
+      );
+    });
+
+    it('should go from DOCUMENT_COUNT to PREVIEW_DATA', () => {
+      expect(getNextStep(MockDataGeneratorStep.DOCUMENT_COUNT)).to.equal(
+        MockDataGeneratorStep.PREVIEW_DATA
+      );
+    });
+
+    it('should go from PREVIEW_DATA to GENERATE_DATA', () => {
+      expect(getNextStep(MockDataGeneratorStep.PREVIEW_DATA)).to.equal(
+        MockDataGeneratorStep.GENERATE_DATA
+      );
+    });
+
+    it('should stay on GENERATE_DATA if already at GENERATE_DATA', () => {
+      expect(getNextStep(MockDataGeneratorStep.GENERATE_DATA)).to.equal(
+        MockDataGeneratorStep.GENERATE_DATA
+      );
+    });
+  });
+
+  describe('getPreviousStep', () => {
+    it('should stay on AI_DISCLAIMER if already at AI_DISCLAIMER', () => {
+      expect(getPreviousStep(MockDataGeneratorStep.AI_DISCLAIMER)).to.equal(
+        MockDataGeneratorStep.AI_DISCLAIMER
+      );
+    });
+
+    it('should go from SCHEMA_CONFIRMATION to AI_DISCLAIMER', () => {
+      expect(
+        getPreviousStep(MockDataGeneratorStep.SCHEMA_CONFIRMATION)
+      ).to.equal(MockDataGeneratorStep.AI_DISCLAIMER);
+    });
+
+    it('should go from SCHEMA_EDITOR to SCHEMA_CONFIRMATION', () => {
+      expect(getPreviousStep(MockDataGeneratorStep.SCHEMA_EDITOR)).to.equal(
+        MockDataGeneratorStep.SCHEMA_CONFIRMATION
+      );
+    });
+
+    it('should go from DOCUMENT_COUNT to SCHEMA_EDITOR', () => {
+      expect(getPreviousStep(MockDataGeneratorStep.DOCUMENT_COUNT)).to.equal(
+        MockDataGeneratorStep.SCHEMA_EDITOR
+      );
+    });
+
+    it('should go from PREVIEW_DATA to DOCUMENT_COUNT', () => {
+      expect(getPreviousStep(MockDataGeneratorStep.PREVIEW_DATA)).to.equal(
+        MockDataGeneratorStep.DOCUMENT_COUNT
+      );
+    });
+
+    it('should go from GENERATE_DATA to PREVIEW_DATA', () => {
+      expect(getPreviousStep(MockDataGeneratorStep.GENERATE_DATA)).to.equal(
+        MockDataGeneratorStep.PREVIEW_DATA
+      );
+    });
+  });
+});
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts b/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
new file mode 100644
index 00000000000..bc273c07dfc
--- /dev/null
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/utils.ts
@@ -0,0 +1,41 @@
+import { MockDataGeneratorStep } from './types';
+
+export const getNextStep = (
+  currentStep: MockDataGeneratorStep
+): MockDataGeneratorStep => {
+  switch (currentStep) {
+    case MockDataGeneratorStep.AI_DISCLAIMER:
+      return MockDataGeneratorStep.SCHEMA_CONFIRMATION;
+    case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
+      return MockDataGeneratorStep.SCHEMA_EDITOR;
+    case MockDataGeneratorStep.SCHEMA_EDITOR:
+      return MockDataGeneratorStep.DOCUMENT_COUNT;
+    case MockDataGeneratorStep.DOCUMENT_COUNT:
+      return MockDataGeneratorStep.PREVIEW_DATA;
+    case MockDataGeneratorStep.PREVIEW_DATA:
+      return MockDataGeneratorStep.GENERATE_DATA;
+    case MockDataGeneratorStep.GENERATE_DATA:
+      // No next step after data generation
+      return currentStep;
+  }
+};
+
+export const getPreviousStep = (
+  currentStep: MockDataGeneratorStep
+): MockDataGeneratorStep => {
+  switch (currentStep) {
+    case MockDataGeneratorStep.AI_DISCLAIMER:
+      // No previous step from AI disclaimer
+      return currentStep;
+    case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
+      return MockDataGeneratorStep.AI_DISCLAIMER;
+    case MockDataGeneratorStep.SCHEMA_EDITOR:
+      return MockDataGeneratorStep.SCHEMA_CONFIRMATION;
+    case MockDataGeneratorStep.DOCUMENT_COUNT:
+      return MockDataGeneratorStep.SCHEMA_EDITOR;
+    case MockDataGeneratorStep.PREVIEW_DATA:
+      return MockDataGeneratorStep.DOCUMENT_COUNT;
+    case MockDataGeneratorStep.GENERATE_DATA:
+      return MockDataGeneratorStep.PREVIEW_DATA;
+  }
+};