feat(compass-collection): Document Count Screen - Mock Data Generator CLOUDP-333856

ncarbon · ncarbon · commit 5905c3194b79 · 2025-09-25T16:30:38.000-04:00
diff --git a/packages/compass-collection/package.json b/packages/compass-collection/package.json
@@ -67,6 +67,7 @@
     "mongodb-collection-model": "^5.34.2",
     "mongodb-ns": "^3.0.1",
     "mongodb-schema": "^12.6.3",
+    "numeral": "^2.0.6",
     "react": "^17.0.2",
     "react-redux": "^8.1.3",
     "redux": "^4.2.1",
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
@@ -7,3 +7,6 @@ export const StepButtonLabelMap = {
   [MockDataGeneratorStep.PREVIEW_DATA]: 'Generate Script',
   [MockDataGeneratorStep.GENERATE_DATA]: 'Done',
 } as const;
+
+export const DEFAULT_DOCUMENT_COUNT = 1000;
+export const MAX_DOCUMENT_COUNT = 100000;
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/document-count-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/document-count-screen.tsx
@@ -7,8 +7,8 @@ import {
 } from '@mongodb-js/compass-components';
 import React, { useMemo } from 'react';
 import { connect } from 'react-redux';
-import { CollectionState } from '../../modules/collection-tab';
-import { SchemaAnalysisState } from '../../schema-analysis-types';
+import type { CollectionState } from '../../modules/collection-tab';
+import type { SchemaAnalysisState } from '../../schema-analysis-types';
 import numeral from 'numeral';
 import { DEFAULT_DOCUMENT_COUNT, MAX_DOCUMENT_COUNT } from './constants';
 
@@ -42,6 +42,15 @@ const formatBytes = (bytes: number) => {
   return numeral(bytes).format(precision + 'b');
 };
 
+type ErrorState =
+  | {
+      state: 'error';
+      message: string;
+    }
+  | {
+      state: 'none';
+    };
+
 interface OwnProps {
   documentCount: number;
   onDocumentCountChange: (documentCount: number) => void;
@@ -56,23 +65,30 @@ const DocumentCountScreen = ({
   onDocumentCountChange,
   schemaAnalysisState,
 }: Props) => {
-  const estimatedDiskSize = useMemo(() => {
-    return schemaAnalysisState.status === 'complete'
-      ? schemaAnalysisState.schemaMetadata.avgDocumentSize * documentCount
-      : 0;
-  }, [schemaAnalysisState, documentCount]);
-
-  const errorState = useMemo(() => {
-    return documentCount < 1 || documentCount > MAX_DOCUMENT_COUNT
-      ? 'error'
-      : 'none';
-  }, [documentCount]);
-
-  const errorMessage = useMemo(() => {
-    return documentCount < 1 || documentCount > MAX_DOCUMENT_COUNT
-      ? 'Document count must be between 1 and 100000'
-      : undefined;
-  }, [documentCount]);
+  const estimatedDiskSize = useMemo(
+    () =>
+      schemaAnalysisState.status === 'complete' &&
+      schemaAnalysisState.schemaMetadata.avgDocumentSize
+        ? formatBytes(
+            schemaAnalysisState.schemaMetadata.avgDocumentSize * documentCount
+          )
+        : 'Not available',
+    [schemaAnalysisState, documentCount]
+  );
+
+  const isOutOfRange = documentCount < 1 || documentCount > MAX_DOCUMENT_COUNT;
+
+  const errorState: ErrorState = useMemo(() => {
+    if (isOutOfRange) {
+      return {
+        state: 'error',
+        message: `Document count must be between 1 and ${MAX_DOCUMENT_COUNT}`,
+      };
+    }
+    return {
+      state: 'none',
+    };
+  }, [isOutOfRange]);
 
   return schemaAnalysisState.status === 'complete' ? (
     <div>
@@ -92,14 +108,14 @@ const DocumentCountScreen = ({
           onChange={(e) => onDocumentCountChange(Number(e.target.value))}
           min={1}
           max={MAX_DOCUMENT_COUNT}
-          state={errorState}
-          errorMessage={errorMessage}
+          state={errorState.state}
+          errorMessage={
+            errorState.state === 'error' ? errorState.message : undefined
+          }
         />
         <div>
           <Body className={boldStyles}>Estimated Disk Size</Body>
-          <Body className={estimatedDiskSizeStyles}>
-            {formatBytes(estimatedDiskSize)}
-          </Body>
+          <Body className={estimatedDiskSizeStyles}>{estimatedDiskSize}</Body>
         </div>
       </div>
     </div>
@@ -109,13 +125,9 @@ const DocumentCountScreen = ({
   );
 };
 
-const mapStateToProps = (state: CollectionState, _ownProps: OwnProps) => {
-  const schemaAnalysisState = state.schemaAnalysis;
-
-  return {
-    schemaAnalysisState,
-  };
-};
+const mapStateToProps = (state: CollectionState) => ({
+  schemaAnalysisState: state.schemaAnalysis,
+});
 
 const ConnectedDocumentCountScreen = connect(
   mapStateToProps,
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
@@ -29,7 +29,11 @@ const defaultSchemaAnalysisState: SchemaAnalysisState = {
     },
   },
   sampleDocument: { name: 'John' },
-  schemaMetadata: { maxNestingDepth: 1, validationRules: null },
+  schemaMetadata: {
+    maxNestingDepth: 1,
+    validationRules: null,
+    avgDocumentSize: undefined,
+  },
 };
 
 describe('MockDataGeneratorModal', () => {
@@ -627,6 +631,85 @@ describe('MockDataGeneratorModal', () => {
     });
   });
 
+  describe('on the document count step', () => {
+    it('displays the correct step title and description', async () => {
+      await renderModal({ currentStep: MockDataGeneratorStep.DOCUMENT_COUNT });
+
+      expect(screen.getByText('Specify Number of Documents to Generate')).to
+        .exist;
+
+      expect(
+        screen.getByText(
+          /Indicate the amount of documents you want to generate below./
+        )
+      ).to.exist;
+      expect(screen.getByText(/Note: We have defaulted to 1000./)).to.exist;
+    });
+
+    it('displays the default document count when the user does not enter a document count', async () => {
+      await renderModal({ currentStep: MockDataGeneratorStep.DOCUMENT_COUNT });
+
+      expect(
+        screen.getByLabelText('Documents to generate in current collection')
+      ).to.have.value('1000');
+    });
+
+    it('disables the Next button and shows an error message when the document count is greater than 100000', async () => {
+      await renderModal({ currentStep: MockDataGeneratorStep.DOCUMENT_COUNT });
+
+      userEvent.type(
+        screen.getByLabelText('Documents to generate in current collection'),
+        '100001'
+      );
+
+      expect(screen.getByText('Document count must be between 1 and 100000')).to
+        .exist;
+      expect(
+        screen.getByTestId('next-step-button').getAttribute('aria-disabled')
+      ).to.equal('true');
+    });
+
+    it('displays "Not available" when the avgDocumentSize is undefined', async () => {
+      await renderModal({
+        currentStep: MockDataGeneratorStep.DOCUMENT_COUNT,
+        schemaAnalysis: {
+          ...defaultSchemaAnalysisState,
+          schemaMetadata: {
+            ...defaultSchemaAnalysisState.schemaMetadata,
+            avgDocumentSize: undefined,
+          },
+        },
+      });
+
+      expect(screen.getByText('Estimated Disk Size')).to.exist;
+      expect(screen.getByText('Not available')).to.exist;
+    });
+
+    it('displays the correct estimated disk size when a valid document count is entered', async () => {
+      await renderModal({
+        currentStep: MockDataGeneratorStep.DOCUMENT_COUNT,
+        schemaAnalysis: {
+          ...defaultSchemaAnalysisState,
+          schemaMetadata: {
+            ...defaultSchemaAnalysisState.schemaMetadata,
+            avgDocumentSize: 100, // 100 bytes
+          },
+        },
+      });
+
+      expect(screen.getByText('Estimated Disk Size')).to.exist;
+      const documentCountInput = screen.getByLabelText(
+        'Documents to generate in current collection'
+      );
+      userEvent.clear(documentCountInput);
+      userEvent.type(documentCountInput, '1000');
+      expect(screen.getByText('100.0KB')).to.exist;
+      userEvent.clear(documentCountInput);
+      userEvent.type(documentCountInput, '2000');
+      expect(screen.getByText('200.0KB')).to.exist;
+    });
+  });
+
   describe('on the generate data step', () => {
     it('enables the Back button', async () => {
       await renderModal({ currentStep: MockDataGeneratorStep.GENERATE_DATA });
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
@@ -14,7 +14,11 @@ import {
 } from '@mongodb-js/compass-components';
 
 import { type MockDataGeneratorState, MockDataGeneratorStep } from './types';
-import { StepButtonLabelMap } from './constants';
+import {
+  DEFAULT_DOCUMENT_COUNT,
+  MAX_DOCUMENT_COUNT,
+  StepButtonLabelMap,
+} from './constants';
 import type { CollectionState } from '../../modules/collection-tab';
 import {
   mockDataGeneratorModalClosed,
@@ -25,6 +29,7 @@ import {
 import RawSchemaConfirmationScreen from './raw-schema-confirmation-screen';
 import FakerSchemaEditorScreen from './faker-schema-editor-screen';
 import ScriptScreen from './script-screen';
+import DocumentCountScreen from './document-count-screen';
 
 const footerStyles = css`
   flex-direction: row;
@@ -65,6 +70,9 @@ const MockDataGeneratorModal = ({
 }: Props) => {
   const [isSchemaConfirmed, setIsSchemaConfirmed] =
     React.useState<boolean>(false);
+  const [documentCount, setDocumentCount] = React.useState<number>(
+    DEFAULT_DOCUMENT_COUNT
+  );
 
   const modalBodyContent = useMemo(() => {
     switch (currentStep) {
@@ -79,16 +87,32 @@ const MockDataGeneratorModal = ({
           />
         );
       case MockDataGeneratorStep.DOCUMENT_COUNT:
-        return <></>; // TODO: CLOUDP-333856
+        return (
+          <DocumentCountScreen
+            documentCount={documentCount}
+            onDocumentCountChange={setDocumentCount}
+          />
+        );
       case MockDataGeneratorStep.PREVIEW_DATA:
         return <></>; // TODO: CLOUDP-333857
       case MockDataGeneratorStep.GENERATE_DATA:
         return <ScriptScreen />;
     }
-  }, [currentStep, fakerSchemaGenerationState, isSchemaConfirmed]);
+  }, [
+    currentStep,
+    fakerSchemaGenerationState,
+    isSchemaConfirmed,
+    documentCount,
+    setDocumentCount,
+  ]);
 
   const isNextButtonDisabled =
-    currentStep === MockDataGeneratorStep.SCHEMA_EDITOR && !isSchemaConfirmed;
+    (currentStep === MockDataGeneratorStep.SCHEMA_EDITOR &&
+      !isSchemaConfirmed) ||
+    (currentStep === MockDataGeneratorStep.DOCUMENT_COUNT &&
+      documentCount < 1) ||
+    (currentStep === MockDataGeneratorStep.DOCUMENT_COUNT &&
+      documentCount > MAX_DOCUMENT_COUNT);
 
   const handleNextClick = () => {
     if (currentStep === MockDataGeneratorStep.GENERATE_DATA) {
diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts
@@ -33,6 +33,7 @@ import {
   processSchema,
   ProcessSchemaUnsupportedStateError,
 } from '../transform-schema-to-field-info';
+import type { Collection } from '@mongodb-js/compass-app-stores/provider';
 import type { Document, MongoError } from 'mongodb';
 import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 import type {
@@ -94,6 +95,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     connectionInfoRef: ConnectionInfoRef;
     fakerSchemaGenerationAbortControllerRef: { current?: AbortController };
     schemaAnalysisAbortControllerRef: { current?: AbortController };
+    collection: Collection;
   },
   A
 >;
@@ -147,6 +149,7 @@ interface SchemaAnalysisFinishedAction {
   schemaMetadata: {
     maxNestingDepth: number;
     validationRules: Document | null;
+    avgDocumentSize: number | undefined;
   };
 }
 
@@ -561,7 +564,13 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
   return async (
     dispatch,
     getState,
-    { dataService, preferences, logger, schemaAnalysisAbortControllerRef }
+    {
+      dataService,
+      preferences,
+      logger,
+      schemaAnalysisAbortControllerRef,
+      collection: collectionModel,
+    }
   ) => {
     const { schemaAnalysis, namespace } = getState();
     const analysisStatus = schemaAnalysis.status;
@@ -638,6 +647,7 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
       const schemaMetadata = {
         maxNestingDepth,
         validationRules,
+        avgDocumentSize: collectionModel.avg_document_size,
       };
 
       // Final check before dispatching results
diff --git a/packages/compass-collection/src/schema-analysis-types.ts b/packages/compass-collection/src/schema-analysis-types.ts
@@ -58,6 +58,7 @@ export type SchemaAnalysisCompletedState = {
   schemaMetadata: {
     maxNestingDepth: number;
     validationRules: Document | null;
+    avgDocumentSize: number | undefined;
   };
 };
 
diff --git a/packages/compass-collection/src/stores/collection-tab.ts b/packages/compass-collection/src/stores/collection-tab.ts
@@ -109,6 +109,7 @@ export function activatePlugin(
     applyMiddleware(
       thunk.withExtraArgument({
         dataService,
+        collection: collectionModel,
         atlasAiService,
         workspaces,
         localAppRegistry,
