From e3b663093df417b94b700b27f14fa42030a12486 Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Wed, 5 Nov 2025 12:40:48 -0500 Subject: [PATCH 01/10] Defaults for unrecognized (eg. _id) --- packages/compass-collection/src/modules/collection-tab.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 09c243f18d1..7917c9f8e0b 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -45,6 +45,7 @@ import type { import { DEFAULT_DOCUMENT_COUNT } from '../components/mock-data-generator-modal/constants'; import { isValidFakerMethod } from '../components/mock-data-generator-modal/utils'; +import { getDefaultFakerMethod } from '../components/mock-data-generator-modal/script-generation-utils'; const DEFAULT_SAMPLE_SIZE = 100; @@ -918,7 +919,7 @@ const validateFakerSchema = ( ); result[fieldPath] = { mongoType: fakerMapping.mongoType, - fakerMethod: UNRECOGNIZED_FAKER_METHOD, + fakerMethod: getDefaultFakerMethod(fakerMapping.mongoType), fakerArgs: [], probability: fakerMapping.probability, }; @@ -927,7 +928,7 @@ const validateFakerSchema = ( // Field not mapped by LLM - add default result[fieldPath] = { mongoType: inputSchema[fieldPath].type, - fakerMethod: UNRECOGNIZED_FAKER_METHOD, + fakerMethod: getDefaultFakerMethod(inputSchema[fieldPath].type), fakerArgs: [], probability: inputSchema[fieldPath].probability, }; From 53dafca522f2ba60f427177f4525c54441537724 Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Wed, 5 Nov 2025 12:54:45 -0500 Subject: [PATCH 02/10] Test --- .../mock-data-generator-modal.spec.tsx | 87 +------------------ 1 file changed, 1 insertion(+), 86 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 071d1960174..3d2f30e7578 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 @@ -6,7 +6,6 @@ import { renderWithActiveConnection, waitFor, userEvent, - waitForElementToBeRemoved, } from '@mongodb-js/testing-library-compass'; import { Provider } from 'react-redux'; import { createStore, applyMiddleware } from 'redux'; @@ -475,23 +474,14 @@ describe('MockDataGeneratorModal', () => { userEvent.click(screen.getByText('email')); expect(screen.getByText('email')).to.exist; expect(screen.getByLabelText('JSON Type')).to.have.value('String'); - // the "email" field should have a warning banner since the faker method is invalid expect(screen.getByLabelText('Faker Function')).to.have.value( - 'Unrecognized' + 'lorem.word' ); - expect( - screen.getByText( - 'Please select a function or we will default fill this field with the string "Unrecognized"' - ) - ).to.exist; // select the "username" field userEvent.click(screen.getByText('username')); expect(screen.getByText('username')).to.exist; expect(screen.getByLabelText('JSON Type')).to.have.value('String'); - expect(screen.getByLabelText('Faker Function')).to.have.value( - 'Unrecognized' - ); }); it('does not show any fields that are not in the input schema', async () => { @@ -532,81 +522,6 @@ describe('MockDataGeneratorModal', () => { expect(screen.queryByText('email')).to.not.exist; }); - it('shows unmapped fields as "Unrecognized"', async () => { - const mockServices = createMockServices(); - mockServices.atlasAiService.getMockDataSchema = () => - Promise.resolve({ - fields: [ - { - fieldPath: 'name', - mongoType: 'String', - fakerMethod: 'person.firstName', - fakerArgs: [], - isArray: false, - probability: 1.0, - }, - { - fieldPath: 'age', - mongoType: 'Int32', - fakerMethod: 'number.int', - fakerArgs: [], - isArray: false, - probability: 1.0, - }, - ], - }); - - await renderModal({ - mockServices, - schemaAnalysis: { - ...defaultSchemaAnalysisState, - processedSchema: { - name: { - type: 'String', - probability: 1.0, - }, - age: { - type: 'Int32', - probability: 1.0, - }, - type: { - type: 'String', - probability: 1.0, - sampleValues: ['cat', 'dog'], - }, - }, - sampleDocument: { name: 'Peaches', age: 10, type: 'cat' }, - }, - }); - - // advance to the schema editor step - userEvent.click(screen.getByText('Confirm')); - await waitForElementToBeRemoved(() => - screen.queryByTestId('faker-schema-editor-loader') - ); - - // select the "name" field - userEvent.click(screen.getByText('name')); - expect(screen.getByLabelText('JSON Type')).to.have.value('String'); - expect(screen.getByLabelText('Faker Function')).to.have.value( - 'person.firstName' - ); - - // select the "age" field - userEvent.click(screen.getByText('age')); - expect(screen.getByLabelText('JSON Type')).to.have.value('Int32'); - expect(screen.getByLabelText('Faker Function')).to.have.value( - 'number.int' - ); - - // select the "type" field - userEvent.click(screen.getByText('type')); - expect(screen.getByLabelText('JSON Type')).to.have.value('String'); - expect(screen.getByLabelText('Faker Function')).to.have.value( - 'Unrecognized' - ); - }); - it('displays preview of the faker call without args when the args are invalid', async () => { const largeLengthArgs = Array.from({ length: 11 }, () => 'testArg'); const mockServices = createMockServices(); From 161c74b7254b988d681445bed4c61ced9f187090 Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Thu, 6 Nov 2025 13:44:02 -0500 Subject: [PATCH 03/10] WIP --- .../raw-schema-confirmation-screen.tsx | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx index 506e74178b2..f44f775fe38 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx @@ -11,9 +11,11 @@ import { DocumentList, useDarkMode, cx, + Link, } from '@mongodb-js/compass-components'; import { usePreference } from 'compass-preferences-model/provider'; +import { useConnectionInfo } from '@mongodb-js/compass-connections/provider'; import toSimplifiedFieldInfo from './to-simplified-field-info'; import type { CollectionState } from '../../modules/collection-tab'; import type { SchemaAnalysisState } from '../../schema-analysis-types'; @@ -44,6 +46,11 @@ const descriptionStyles = css({ marginBottom: spacing[200], }); +const projectSettingsInfoStyles = css({ + marginBottom: spacing[400], + color: palette.gray.dark1, +}); + const errorBannerStyles = css({ marginTop: spacing[400], marginBottom: spacing[400], @@ -61,6 +68,7 @@ const RawSchemaConfirmationScreen = ({ 'enableGenAISampleDocumentPassing' ); const isDarkMode = useDarkMode(); + const connectionInfo = useConnectionInfo(); const subtitleText = enableSampleDocumentPassing ? 'Sample Documents Collected' @@ -70,6 +78,19 @@ const RawSchemaConfirmationScreen = ({ ? 'A sample of documents from your collection will be sent to an LLM for processing.' : 'We have identified the following schema from your documents. This schema will be sent to an LLM for processing.'; + const projectId = connectionInfo.atlasMetadata?.projectId; + const projectSettingsLink = projectId ? ( + + Project Settings + + ) : ( + 'Project Settings' + ); + return (
{schemaAnalysis.status === 'complete' ? ( @@ -78,6 +99,11 @@ const RawSchemaConfirmationScreen = ({ {subtitleText} {descriptionText} + + To improve mock data quality, Project Owners can enable sending + sample field values to the AI model in {projectSettingsLink}. + Refresh Data Explorer for changes to take effect. + {fakerSchemaGenerationStatus === 'error' && ( Date: Thu, 6 Nov 2025 14:17:40 -0500 Subject: [PATCH 04/10] Persist selections --- .../faker-schema-editor-screen.tsx | 121 ++++++++++-------- .../mock-data-generator-modal.spec.tsx | 58 +++++++++ 2 files changed, 123 insertions(+), 56 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 744a65efb57..5669d364199 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -11,16 +11,17 @@ import { } from '@mongodb-js/compass-components'; import React from 'react'; import { connect } from 'react-redux'; +import type { Dispatch } from 'redux'; import FieldSelector from './schema-field-selector'; import FakerMappingSelector from './faker-mapping-selector'; import { getDefaultFakerMethod } from './script-generation-utils'; -import type { - FakerSchema, - FakerFieldMapping, - MockDataGeneratorState, -} from './types'; +import type { FakerSchema, MockDataGeneratorState } from './types'; import type { MongoDBFieldType } from '../../schema-analysis-types'; import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; +import { + fakerFieldTypeChanged, + fakerFieldMethodChanged, +} from '../../modules/collection-tab'; const containerStyles = css({ display: 'flex', @@ -58,37 +59,29 @@ const schemaEditorLoaderStyles = css({ const FakerSchemaEditorContent = ({ fakerSchema, + originalLlmResponse, onSchemaConfirmed, + onFieldTypeChanged, + onFieldMethodChanged, }: { fakerSchema: FakerSchema; + originalLlmResponse: FakerSchema; onSchemaConfirmed: () => void; + onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; + onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; }) => { const track = useTelemetry(); - const [fakerSchemaFormValues, setFakerSchemaFormValues] = - React.useState(fakerSchema); - - // Store original LLM mappings to restore when reselecting original methods - const originalLlmMappings = React.useRef>( - Object.fromEntries( - Object.entries(fakerSchema).map(([field, mapping]) => [ - field, - { - ...mapping, - }, - ]) - ) - ); - const fieldPaths = Object.keys(fakerSchemaFormValues); + const fieldPaths = Object.keys(fakerSchema); const [activeField, setActiveField] = React.useState(fieldPaths[0]); - const activeJsonType = fakerSchemaFormValues[activeField]?.mongoType; - const activeFakerFunction = fakerSchemaFormValues[activeField]?.fakerMethod; - const activeFakerArgs = fakerSchemaFormValues[activeField]?.fakerArgs; + const activeJsonType = fakerSchema[activeField]?.mongoType; + const activeFakerFunction = fakerSchema[activeField]?.fakerMethod; + const activeFakerArgs = fakerSchema[activeField]?.fakerArgs; const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => { - const currentMapping = fakerSchemaFormValues[activeField]; - const originalLlmMapping = originalLlmMappings.current[activeField]; + const currentMapping = fakerSchema[activeField]; + const originalLlmMapping = originalLlmResponse[activeField]; if (currentMapping) { const previousJsonType = currentMapping.mongoType; @@ -97,16 +90,20 @@ const FakerSchemaEditorContent = ({ const isSwitchingToOriginalType = originalLlmMapping && newJsonType === originalLlmMapping.mongoType; - const newMapping = isSwitchingToOriginalType - ? { ...originalLlmMapping } - : { - ...currentMapping, - mongoType: newJsonType, - fakerMethod: getDefaultFakerMethod(newJsonType), - fakerArgs: [], - }; - - const newFakerMethod = newMapping.fakerMethod; + if (isSwitchingToOriginalType) { + // Restore original LLM mapping + onFieldTypeChanged(activeField, originalLlmMapping.mongoType); + onFieldMethodChanged(activeField, originalLlmMapping.fakerMethod); + } else { + // Use default faker method for new type + const newFakerMethod = getDefaultFakerMethod(newJsonType); + onFieldTypeChanged(activeField, newJsonType); + onFieldMethodChanged(activeField, newFakerMethod); + } + + const newFakerMethod = isSwitchingToOriginalType + ? originalLlmMapping.fakerMethod + : getDefaultFakerMethod(newJsonType); track('Mock Data JSON Type Changed', { field_name: activeField, @@ -115,17 +112,12 @@ const FakerSchemaEditorContent = ({ previous_faker_method: previousFakerMethod, new_faker_method: newFakerMethod, }); - - setFakerSchemaFormValues({ - ...fakerSchemaFormValues, - [activeField]: newMapping, - }); } }; const onFakerFunctionSelect = (newFakerFunction: string) => { - const currentMapping = fakerSchemaFormValues[activeField]; - const originalLlmMapping = originalLlmMappings.current[activeField]; + const currentMapping = fakerSchema[activeField]; + const originalLlmMapping = originalLlmResponse[activeField]; if (currentMapping) { const previousFakerMethod = currentMapping.fakerMethod; @@ -135,13 +127,14 @@ const FakerSchemaEditorContent = ({ currentMapping.mongoType === originalLlmMapping.mongoType && newFakerFunction === originalLlmMapping.fakerMethod; - const newMapping = isSwitchingToLlmSuggestion - ? { ...originalLlmMapping } - : { - ...currentMapping, - fakerMethod: newFakerFunction, - fakerArgs: [], - }; + if (isSwitchingToLlmSuggestion) { + // Restore original LLM mapping completely + onFieldTypeChanged(activeField, originalLlmMapping.mongoType); + onFieldMethodChanged(activeField, originalLlmMapping.fakerMethod); + } else { + // Just update the faker method + onFieldMethodChanged(activeField, newFakerFunction); + } track('Mock Data Faker Method Changed', { field_name: activeField, @@ -149,11 +142,6 @@ const FakerSchemaEditorContent = ({ previous_faker_method: previousFakerMethod, new_faker_method: newFakerFunction, }); - - setFakerSchemaFormValues({ - ...fakerSchemaFormValues, - [activeField]: newMapping, - }); } }; @@ -164,7 +152,7 @@ const FakerSchemaEditorContent = ({ activeField={activeField} fields={fieldPaths} onFieldSelect={setActiveField} - fakerSchema={fakerSchemaFormValues} + fakerSchema={fakerSchema} /> {activeJsonType && activeFakerFunction && ( void; fakerSchemaGenerationState: MockDataGeneratorState; + onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; + onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; }) => { return (
@@ -220,11 +212,28 @@ const FakerSchemaEditorScreen = ({ {fakerSchemaGenerationState.status === 'completed' && ( )}
); }; -export default connect()(FakerSchemaEditorScreen); +const mapStateToProps = () => ({ + // No additional state needed - component already receives fakerSchemaGenerationState as prop +}); + +const mapDispatchToProps = (dispatch: Dispatch) => ({ + onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => + dispatch(fakerFieldTypeChanged(fieldPath, mongoType)), + onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => + dispatch(fakerFieldMethodChanged(fieldPath, fakerMethod)), +}); + +export default connect( + mapStateToProps, + mapDispatchToProps +)(FakerSchemaEditorScreen); 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 3d2f30e7578..0df1f742f7c 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 @@ -733,6 +733,64 @@ describe('MockDataGeneratorModal', () => { ); }); }); + + it('persists user modifications to faker mappings when navigating back to the screen', async () => { + await renderModal({ + mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, + }); + + userEvent.click(screen.getByText('Confirm')); + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + + // Change the JSON type from String to Number + const jsonTypeSelect = screen.getByLabelText('JSON Type'); + userEvent.click(jsonTypeSelect); + const numberOption = await screen.findByRole('option', { + name: 'Number', + }); + userEvent.click(numberOption); + + await waitFor(() => { + expect(screen.getByLabelText('JSON Type')).to.have.value('Number'); + }); + + const fakerMethodSelect = screen.getByLabelText('Faker Function'); + userEvent.click(fakerMethodSelect); + const floatOption = await screen.findByRole('option', { + name: 'number.float', + }); + userEvent.click(floatOption); + + await waitFor(() => { + expect(screen.getByLabelText('Faker Function')).to.have.value( + 'number.float' + ); + }); + + // Advance to document count step + userEvent.click(screen.getByText('Confirm mappings')); + await waitFor(() => { + expect(screen.getByText('Specify Number of Documents to Generate')).to + .exist; + }); + + // Go back to schema editor + userEvent.click(screen.getByText('Back')); + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + + // Verify both selections persisted + await waitFor(() => { + expect(screen.getByLabelText('JSON Type')).to.have.value('Number'); + expect(screen.getByLabelText('Faker Function')).to.have.value( + 'number.float' + ); + }); + }); }); describe('on the document count step', () => { From 8124613cb765cd0b8c2dcd06eea48b35bd8d5592 Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Thu, 6 Nov 2025 18:02:29 -0500 Subject: [PATCH 05/10] WIP --- .../faker-mapping-selector.spec.tsx | 49 +++++++++++++++ .../faker-mapping-selector.tsx | 11 ++-- .../faker-schema-editor-screen.tsx | 53 ++++++++++------ .../mock-data-generator-modal.spec.tsx | 62 +++++++++++++++++++ .../src/modules/collection-tab.ts | 45 ++++++++++++++ 5 files changed, 196 insertions(+), 24 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx index edae8dedc29..9536046f93a 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx @@ -157,4 +157,53 @@ describe('FakerMappingSelector', () => { ) ).to.not.exist; }); + + it('should always include the original LLM faker method in the dropdown', () => { + const originalLlmMethod = 'custom.llmMethod'; + + render( + + ); + + const fakerFunctionSelect = screen.getByLabelText('Faker Function'); + userEvent.click(fakerFunctionSelect); + + // Should include the original LLM method even though it's not in MONGO_TYPE_TO_FAKER_METHODS + expect(screen.getByRole('option', { name: originalLlmMethod })).to.exist; + + // Should also include standard methods for String type + expect(screen.getByRole('option', { name: 'lorem.word' })).to.exist; + expect(screen.getByRole('option', { name: 'lorem.sentence' })).to.exist; + }); + + it('should not duplicate the original LLM method if it is already in the standard methods', () => { + const originalLlmMethod = 'lorem.word'; + + render( + + ); + + const fakerFunctionSelect = screen.getByLabelText('Faker Function'); + userEvent.click(fakerFunctionSelect); + + // Should only have one instance of 'lorem.word' + const loremWordOptions = screen.getAllByRole('option', { + name: 'lorem.word', + }); + expect(loremWordOptions).to.have.length(1); + }); }); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx index aa08248262b..494a256c516 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx @@ -67,6 +67,7 @@ interface Props { onJsonTypeSelect: (jsonType: MongoDBFieldType) => void; activeFakerArgs: FakerArg[]; onFakerFunctionSelect: (fakerFunction: string) => void; + originalLlmFakerMethod?: string; } const FakerMappingSelector = ({ @@ -75,16 +76,18 @@ const FakerMappingSelector = ({ activeFakerArgs, onJsonTypeSelect, onFakerFunctionSelect, + originalLlmFakerMethod, }: Props) => { const fakerMethodOptions = useMemo(() => { const methods = MONGO_TYPE_TO_FAKER_METHODS[activeJsonType] || []; - if (methods.includes(activeFakerFunction)) { - return methods; + // Include original LLM method if it's not already in the list of methods + if (originalLlmFakerMethod && !methods.includes(originalLlmFakerMethod)) { + return [originalLlmFakerMethod, ...methods]; } - return [activeFakerFunction, ...methods]; - }, [activeJsonType, activeFakerFunction]); + return methods; + }, [activeJsonType, originalLlmFakerMethod]); return (
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 5669d364199..3a67c18dbb2 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -15,12 +15,17 @@ import type { Dispatch } from 'redux'; import FieldSelector from './schema-field-selector'; import FakerMappingSelector from './faker-mapping-selector'; import { getDefaultFakerMethod } from './script-generation-utils'; -import type { FakerSchema, MockDataGeneratorState } from './types'; +import type { + FakerSchema, + MockDataGeneratorState, + FakerFieldMapping, +} from './types'; import type { MongoDBFieldType } from '../../schema-analysis-types'; import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; import { fakerFieldTypeChanged, fakerFieldMethodChanged, + fakerFieldMappingRestored, } from '../../modules/collection-tab'; const containerStyles = css({ @@ -63,12 +68,17 @@ const FakerSchemaEditorContent = ({ onSchemaConfirmed, onFieldTypeChanged, onFieldMethodChanged, + onFieldMappingRestored, }: { fakerSchema: FakerSchema; originalLlmResponse: FakerSchema; onSchemaConfirmed: () => void; onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; + onFieldMappingRestored: ( + fieldPath: string, + mapping: FakerFieldMapping + ) => void; }) => { const track = useTelemetry(); @@ -90,21 +100,19 @@ const FakerSchemaEditorContent = ({ const isSwitchingToOriginalType = originalLlmMapping && newJsonType === originalLlmMapping.mongoType; + const newFakerMethod = isSwitchingToOriginalType + ? originalLlmMapping.fakerMethod + : getDefaultFakerMethod(newJsonType); + if (isSwitchingToOriginalType) { // Restore original LLM mapping - onFieldTypeChanged(activeField, originalLlmMapping.mongoType); - onFieldMethodChanged(activeField, originalLlmMapping.fakerMethod); + onFieldMappingRestored(activeField, originalLlmMapping); } else { // Use default faker method for new type - const newFakerMethod = getDefaultFakerMethod(newJsonType); onFieldTypeChanged(activeField, newJsonType); onFieldMethodChanged(activeField, newFakerMethod); } - const newFakerMethod = isSwitchingToOriginalType - ? originalLlmMapping.fakerMethod - : getDefaultFakerMethod(newJsonType); - track('Mock Data JSON Type Changed', { field_name: activeField, previous_json_type: previousJsonType, @@ -128,11 +136,10 @@ const FakerSchemaEditorContent = ({ newFakerFunction === originalLlmMapping.fakerMethod; if (isSwitchingToLlmSuggestion) { - // Restore original LLM mapping completely - onFieldTypeChanged(activeField, originalLlmMapping.mongoType); - onFieldMethodChanged(activeField, originalLlmMapping.fakerMethod); + // Restore original LLM mapping + onFieldMappingRestored(activeField, originalLlmMapping); } else { - // Just update the faker method + // Update the faker method onFieldMethodChanged(activeField, newFakerFunction); } @@ -161,6 +168,11 @@ const FakerSchemaEditorContent = ({ activeFakerArgs={activeFakerArgs} onJsonTypeSelect={onJsonTypeSelect} onFakerFunctionSelect={onFakerFunctionSelect} + originalLlmFakerMethod={ + originalLlmResponse[activeField]?.mongoType === activeJsonType + ? originalLlmResponse[activeField]?.fakerMethod + : undefined + } /> )}
@@ -181,11 +193,16 @@ const FakerSchemaEditorScreen = ({ fakerSchemaGenerationState, onFieldTypeChanged, onFieldMethodChanged, + onFieldMappingRestored, }: { onSchemaConfirmed: () => void; fakerSchemaGenerationState: MockDataGeneratorState; onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; + onFieldMappingRestored: ( + fieldPath: string, + mapping: FakerFieldMapping + ) => void; }) => { return (
@@ -216,24 +233,20 @@ const FakerSchemaEditorScreen = ({ onSchemaConfirmed={onSchemaConfirmed} onFieldTypeChanged={onFieldTypeChanged} onFieldMethodChanged={onFieldMethodChanged} + onFieldMappingRestored={onFieldMappingRestored} /> )}
); }; -const mapStateToProps = () => ({ - // No additional state needed - component already receives fakerSchemaGenerationState as prop -}); - const mapDispatchToProps = (dispatch: Dispatch) => ({ onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => dispatch(fakerFieldTypeChanged(fieldPath, mongoType)), onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => dispatch(fakerFieldMethodChanged(fieldPath, fakerMethod)), + onFieldMappingRestored: (fieldPath: string, mapping: FakerFieldMapping) => + dispatch(fakerFieldMappingRestored(fieldPath, mapping)), }); -export default connect( - mapStateToProps, - mapDispatchToProps -)(FakerSchemaEditorScreen); +export default connect(() => ({}), mapDispatchToProps)(FakerSchemaEditorScreen); 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 0df1f742f7c..32647dd7d08 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 @@ -791,6 +791,68 @@ describe('MockDataGeneratorModal', () => { ); }); }); + + it('preserves original fakerArgs when selecting original LLM method', async () => { + // Mock response with fakerArgs + const mockServicesWithArgs = { + ...mockServicesWithMockDataResponse, + atlasAiService: { + getMockDataSchema: sinon.stub().resolves({ + fields: [ + { + fieldPath: 'name', + mongoType: 'String', + fakerMethod: 'person.firstName', + fakerArgs: [{ json: '{"locale":"en"}' }], + probability: 0.8, + }, + ], + }), + }, + }; + + await renderModal({ + mockServices: mockServicesWithArgs, + schemaAnalysis: mockSchemaAnalysis, + }); + + userEvent.click(screen.getByText('Confirm')); + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + + // Change to a different faker method (this resets fakerArgs) + const fakerMethodSelect = screen.getByLabelText('Faker Function'); + userEvent.click(fakerMethodSelect); + const wordOption = await screen.findByRole('option', { + name: 'lorem.word', + }); + userEvent.click(wordOption); + + await waitFor(() => { + expect(screen.getByLabelText('Faker Function')).to.have.value( + 'lorem.word' + ); + }); + + // Select the original LLM method again + userEvent.click(fakerMethodSelect); + const firstNameOption = await screen.findByRole('option', { + name: 'person.firstName', + }); + userEvent.click(firstNameOption); + + await waitFor(() => { + expect(screen.getByLabelText('Faker Function')).to.have.value( + 'person.firstName' + ); + }); + + const preview = screen.getByTestId('faker-function-call-preview'); + expect(preview.textContent).to.include( + 'faker.person.firstName({"locale":"en"})' + ); + }); }); describe('on the document count step', () => { diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 7917c9f8e0b..4c7bcdc8f8a 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -40,6 +40,7 @@ import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/t import type { LlmFakerMapping, FakerSchema, + FakerFieldMapping, MockDataGeneratorState, } from '../components/mock-data-generator-modal/types'; import { DEFAULT_DOCUMENT_COUNT } from '../components/mock-data-generator-modal/constants'; @@ -134,6 +135,7 @@ export enum CollectionActions { FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed', FakerFieldTypeChanged = 'compass-collection/FakerFieldTypeChanged', FakerFieldMethodChanged = 'compass-collection/FakerFieldMethodChanged', + FakerFieldMappingRestored = 'compass-collection/FakerFieldMappingRestored', } interface CollectionMetadataFetchedAction { @@ -220,6 +222,12 @@ export interface FakerFieldMethodChangedAction { fakerMethod: string; } +export interface FakerFieldMappingRestoredAction { + type: CollectionActions.FakerFieldMappingRestored; + fieldPath: string; + mapping: FakerFieldMapping; +} + const reducer: Reducer = ( state = { // TODO(COMPASS-7782): use hook to get the workspace tab id instead @@ -556,6 +564,7 @@ const reducer: Reducer = ( [fieldPath]: { ...currentMapping, mongoType, + fakerArgs: [], // Reset args when type changes }, }, }, @@ -589,12 +598,37 @@ const reducer: Reducer = ( [fieldPath]: { ...currentMapping, fakerMethod, + fakerArgs: [], // Reset args when method changes }, }, }, }; } + if ( + isAction( + action, + CollectionActions.FakerFieldMappingRestored + ) + ) { + if (state.fakerSchemaGeneration.status !== 'completed') { + return state; + } + + const { fieldPath, mapping } = action; + + return { + ...state, + fakerSchemaGeneration: { + ...state.fakerSchemaGeneration, + editedFakerSchema: { + ...state.fakerSchemaGeneration.editedFakerSchema, + [fieldPath]: mapping, + }, + }, + }; + } + return state; }; @@ -667,6 +701,17 @@ export const fakerFieldMethodChanged = ( }; }; +export const fakerFieldMappingRestored = ( + fieldPath: string, + mapping: FakerFieldMapping +): FakerFieldMappingRestoredAction => { + return { + type: CollectionActions.FakerFieldMappingRestored, + fieldPath, + mapping, + }; +}; + export const selectTab = ( tabName: CollectionSubtab ): CollectionThunkAction => { From 4f9a639b08673fb3891131eafa92b18e586fd8fb Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Thu, 6 Nov 2025 18:50:46 -0500 Subject: [PATCH 06/10] Null instead of empty function --- .../mock-data-generator-modal/faker-schema-editor-screen.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 3a67c18dbb2..36d2121163a 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -249,4 +249,4 @@ const mapDispatchToProps = (dispatch: Dispatch) => ({ dispatch(fakerFieldMappingRestored(fieldPath, mapping)), }); -export default connect(() => ({}), mapDispatchToProps)(FakerSchemaEditorScreen); +export default connect(null, mapDispatchToProps)(FakerSchemaEditorScreen); From 171ae4aa7108316383cb5b03200dff80a96a2d91 Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Mon, 10 Nov 2025 15:53:58 -0500 Subject: [PATCH 07/10] Remove code for unrecognized --- .../faker-mapping-selector.spec.tsx | 37 ------------------- .../faker-mapping-selector.tsx | 37 ++++++------------- 2 files changed, 11 insertions(+), 63 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx index 9536046f93a..962b5fd8212 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx @@ -9,7 +9,6 @@ import { } from '@mongodb-js/testing-library-compass'; import sinon from 'sinon'; import FakerMappingSelector from './faker-mapping-selector'; -import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab'; import type { MongoDBFieldType } from '../../schema-analysis-types'; import { MONGO_TYPE_TO_FAKER_METHODS, @@ -122,42 +121,6 @@ describe('FakerMappingSelector', () => { ); }); - it('should show warning banner when faker method is unrecognized', () => { - render( - - ); - - expect( - screen.getByText( - /Please select a function or we will default fill this field/ - ) - ).to.exist; - }); - - it('should not show warning banner when faker method is recognized', () => { - render( - - ); - - expect( - screen.queryByText( - /Please select a function or we will default fill this field/ - ) - ).to.not.exist; - }); - it('should always include the original LLM faker method in the dropdown', () => { const originalLlmMethod = 'custom.llmMethod'; diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx index 494a256c516..7307444b212 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx @@ -1,6 +1,4 @@ import { - Banner, - BannerVariant, Body, Code, css, @@ -11,7 +9,6 @@ import { spacing, } from '@mongodb-js/compass-components'; import React, { useMemo } from 'react'; -import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab'; import { MONGO_TYPE_TO_FAKER_METHODS, MongoDBFieldTypeValues, @@ -116,29 +113,17 @@ const FakerMappingSelector = ({ ))} - {activeFakerFunction === UNRECOGNIZED_FAKER_METHOD ? ( - - Please select a function or we will default fill this field with the - string "Unrecognized" - - ) : ( - <> - - - {formatFakerFunctionCallWithArgs( - activeFakerFunction, - activeFakerArgs - )} - - - )} + + + {formatFakerFunctionCallWithArgs(activeFakerFunction, activeFakerArgs)} +
); }; From 12fab075deb1d9aea76857fca6e85fd9e9cbaa2e Mon Sep 17 00:00:00 2001 From: Jacob Samuel Lu Date: Mon, 10 Nov 2025 23:40:32 -0500 Subject: [PATCH 08/10] Remove unrecognized icon --- .../faker-schema-editor-screen.tsx | 1 - .../schema-field-selector.tsx | 17 ----------------- 2 files changed, 18 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 36d2121163a..5462e44cc10 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -159,7 +159,6 @@ const FakerSchemaEditorContent = ({ activeField={activeField} fields={fieldPaths} onFieldSelect={setActiveField} - fakerSchema={fakerSchema} /> {activeJsonType && activeFakerFunction && ( void; fields: Array; - fakerSchema?: FakerSchema; -}; - -const shouldShowUnrecognizedIcon = ( - field: string, - fakerSchema?: FakerSchema -): boolean => { - const mapping = fakerSchema?.[field]; - - return !!mapping && mapping.fakerMethod === UNRECOGNIZED_FAKER_METHOD; }; const FieldSelector: React.FunctionComponent = ({ activeField, fields, onFieldSelect, - fakerSchema, }) => { const darkMode = useDarkMode(); @@ -126,9 +112,6 @@ const FieldSelector: React.FunctionComponent = ({ onClick={() => onFieldSelect(field)} > {field} - {shouldShowUnrecognizedIcon(field, fakerSchema) && ( - - )} ))} From 63aa961b9d7f6b3685a2bc3d6dab58d9ed45654d Mon Sep 17 00:00:00 2001 From: Kyle W Lai Date: Wed, 12 Nov 2025 12:01:47 -0500 Subject: [PATCH 09/10] Handle restoring original LLM mapping in reducer --- .../faker-schema-editor-screen.tsx | 54 +++------------- .../src/modules/collection-tab.ts | 61 +++++-------------- 2 files changed, 22 insertions(+), 93 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 5462e44cc10..54b4a6bd132 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -15,17 +15,12 @@ import type { Dispatch } from 'redux'; import FieldSelector from './schema-field-selector'; import FakerMappingSelector from './faker-mapping-selector'; import { getDefaultFakerMethod } from './script-generation-utils'; -import type { - FakerSchema, - MockDataGeneratorState, - FakerFieldMapping, -} from './types'; +import type { FakerSchema, MockDataGeneratorState } from './types'; import type { MongoDBFieldType } from '../../schema-analysis-types'; import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; import { fakerFieldTypeChanged, fakerFieldMethodChanged, - fakerFieldMappingRestored, } from '../../modules/collection-tab'; const containerStyles = css({ @@ -68,17 +63,12 @@ const FakerSchemaEditorContent = ({ onSchemaConfirmed, onFieldTypeChanged, onFieldMethodChanged, - onFieldMappingRestored, }: { fakerSchema: FakerSchema; originalLlmResponse: FakerSchema; onSchemaConfirmed: () => void; onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; - onFieldMappingRestored: ( - fieldPath: string, - mapping: FakerFieldMapping - ) => void; }) => { const track = useTelemetry(); @@ -97,21 +87,13 @@ const FakerSchemaEditorContent = ({ const previousJsonType = currentMapping.mongoType; const previousFakerMethod = currentMapping.fakerMethod; - const isSwitchingToOriginalType = - originalLlmMapping && newJsonType === originalLlmMapping.mongoType; + const newFakerMethod = + originalLlmMapping && newJsonType === originalLlmMapping.mongoType + ? originalLlmMapping.fakerMethod + : getDefaultFakerMethod(newJsonType); - const newFakerMethod = isSwitchingToOriginalType - ? originalLlmMapping.fakerMethod - : getDefaultFakerMethod(newJsonType); - - if (isSwitchingToOriginalType) { - // Restore original LLM mapping - onFieldMappingRestored(activeField, originalLlmMapping); - } else { - // Use default faker method for new type - onFieldTypeChanged(activeField, newJsonType); - onFieldMethodChanged(activeField, newFakerMethod); - } + onFieldTypeChanged(activeField, newJsonType); + onFieldMethodChanged(activeField, newFakerMethod); track('Mock Data JSON Type Changed', { field_name: activeField, @@ -125,23 +107,11 @@ const FakerSchemaEditorContent = ({ const onFakerFunctionSelect = (newFakerFunction: string) => { const currentMapping = fakerSchema[activeField]; - const originalLlmMapping = originalLlmResponse[activeField]; if (currentMapping) { const previousFakerMethod = currentMapping.fakerMethod; - const isSwitchingToLlmSuggestion = - originalLlmMapping && - currentMapping.mongoType === originalLlmMapping.mongoType && - newFakerFunction === originalLlmMapping.fakerMethod; - - if (isSwitchingToLlmSuggestion) { - // Restore original LLM mapping - onFieldMappingRestored(activeField, originalLlmMapping); - } else { - // Update the faker method - onFieldMethodChanged(activeField, newFakerFunction); - } + onFieldMethodChanged(activeField, newFakerFunction); track('Mock Data Faker Method Changed', { field_name: activeField, @@ -192,16 +162,11 @@ const FakerSchemaEditorScreen = ({ fakerSchemaGenerationState, onFieldTypeChanged, onFieldMethodChanged, - onFieldMappingRestored, }: { onSchemaConfirmed: () => void; fakerSchemaGenerationState: MockDataGeneratorState; onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void; onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void; - onFieldMappingRestored: ( - fieldPath: string, - mapping: FakerFieldMapping - ) => void; }) => { return (
@@ -232,7 +197,6 @@ const FakerSchemaEditorScreen = ({ onSchemaConfirmed={onSchemaConfirmed} onFieldTypeChanged={onFieldTypeChanged} onFieldMethodChanged={onFieldMethodChanged} - onFieldMappingRestored={onFieldMappingRestored} /> )}
@@ -244,8 +208,6 @@ const mapDispatchToProps = (dispatch: Dispatch) => ({ dispatch(fakerFieldTypeChanged(fieldPath, mongoType)), onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => dispatch(fakerFieldMethodChanged(fieldPath, fakerMethod)), - onFieldMappingRestored: (fieldPath: string, mapping: FakerFieldMapping) => - dispatch(fakerFieldMappingRestored(fieldPath, mapping)), }); export default connect(null, mapDispatchToProps)(FakerSchemaEditorScreen); diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 4c7bcdc8f8a..916c0a26d8b 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -40,7 +40,6 @@ import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/t import type { LlmFakerMapping, FakerSchema, - FakerFieldMapping, MockDataGeneratorState, } from '../components/mock-data-generator-modal/types'; import { DEFAULT_DOCUMENT_COUNT } from '../components/mock-data-generator-modal/constants'; @@ -135,7 +134,6 @@ export enum CollectionActions { FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed', FakerFieldTypeChanged = 'compass-collection/FakerFieldTypeChanged', FakerFieldMethodChanged = 'compass-collection/FakerFieldMethodChanged', - FakerFieldMappingRestored = 'compass-collection/FakerFieldMappingRestored', } interface CollectionMetadataFetchedAction { @@ -222,12 +220,6 @@ export interface FakerFieldMethodChangedAction { fakerMethod: string; } -export interface FakerFieldMappingRestoredAction { - type: CollectionActions.FakerFieldMappingRestored; - fieldPath: string; - mapping: FakerFieldMapping; -} - const reducer: Reducer = ( state = { // TODO(COMPASS-7782): use hook to get the workspace tab id instead @@ -584,38 +576,24 @@ const reducer: Reducer = ( const { fieldPath, fakerMethod } = action; const currentMapping = state.fakerSchemaGeneration.editedFakerSchema[fieldPath]; + const originalLlmMapping = + state.fakerSchemaGeneration.originalLlmResponse[fieldPath]; if (!currentMapping) { return state; } - return { - ...state, - fakerSchemaGeneration: { - ...state.fakerSchemaGeneration, - editedFakerSchema: { - ...state.fakerSchemaGeneration.editedFakerSchema, - [fieldPath]: { - ...currentMapping, - fakerMethod, - fakerArgs: [], // Reset args when method changes - }, - }, - }, - }; - } - - if ( - isAction( - action, - CollectionActions.FakerFieldMappingRestored - ) - ) { - if (state.fakerSchemaGeneration.status !== 'completed') { - return state; - } - - const { fieldPath, mapping } = action; + const isRestoringOriginalMapping = + originalLlmMapping && + currentMapping.mongoType === originalLlmMapping.mongoType && + fakerMethod === originalLlmMapping.fakerMethod; + const updatedMapping = isRestoringOriginalMapping + ? originalLlmMapping + : { + ...currentMapping, + fakerMethod, + fakerArgs: [], // Reset args when method changes + }; return { ...state, @@ -623,7 +601,7 @@ const reducer: Reducer = ( ...state.fakerSchemaGeneration, editedFakerSchema: { ...state.fakerSchemaGeneration.editedFakerSchema, - [fieldPath]: mapping, + [fieldPath]: updatedMapping, }, }, }; @@ -701,17 +679,6 @@ export const fakerFieldMethodChanged = ( }; }; -export const fakerFieldMappingRestored = ( - fieldPath: string, - mapping: FakerFieldMapping -): FakerFieldMappingRestoredAction => { - return { - type: CollectionActions.FakerFieldMappingRestored, - fieldPath, - mapping, - }; -}; - export const selectTab = ( tabName: CollectionSubtab ): CollectionThunkAction => { From c26678c289542c24d0f1c0e6aa6caa657a67cb75 Mon Sep 17 00:00:00 2001 From: Kyle W Lai Date: Thu, 13 Nov 2025 14:23:01 -0500 Subject: [PATCH 10/10] Fix failing test --- .../mock-data-generator-modal.spec.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 688709a48f5..55f13900643 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 @@ -771,7 +771,7 @@ describe('MockDataGeneratorModal', () => { }); // Advance to document count step - userEvent.click(screen.getByText('Confirm mappings')); + userEvent.click(screen.getByTestId('next-step-button')); await waitFor(() => { expect(screen.getByText('Specify Number of Documents to Generate')).to .exist;