feat(compass-collection): Mock Data Generator QA Items Batch 1 - CLOUDP-356788

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(
+      <FakerMappingSelector
+        activeJsonType="String"
+        activeFakerFunction="lorem.word"
+        activeFakerArgs={[]}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+        originalLlmFakerMethod={originalLlmMethod}
+      />
+    );
+
+    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(
+      <FakerMappingSelector
+        activeJsonType="String"
+        activeFakerFunction="lorem.sentence"
+        activeFakerArgs={[]}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+        originalLlmFakerMethod={originalLlmMethod}
+      />
+    );
+
+    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);
+  });
 });

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 (
     <div className={fieldMappingSelectorsStyles}>

packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx

@@ -11,16 +11,22 @@ 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,
+  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({
   display: 'flex',
@@ -58,37 +64,34 @@ const schemaEditorLoaderStyles = css({
 
 const FakerSchemaEditorContent = ({
   fakerSchema,
+  originalLlmResponse,
   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();
-  const [fakerSchemaFormValues, setFakerSchemaFormValues] =
-    React.useState<FakerSchema>(fakerSchema);
-
-  // Store original LLM mappings to restore when reselecting original methods
-  const originalLlmMappings = React.useRef<Record<string, FakerFieldMapping>>(
-    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<string>(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 +100,18 @@ const FakerSchemaEditorContent = ({
       const isSwitchingToOriginalType =
         originalLlmMapping && newJsonType === originalLlmMapping.mongoType;
 
-      const newMapping = isSwitchingToOriginalType
-        ? { ...originalLlmMapping }
-        : {
-            ...currentMapping,
-            mongoType: newJsonType,
-            fakerMethod: getDefaultFakerMethod(newJsonType),
-            fakerArgs: [],
-          };
+      const newFakerMethod = isSwitchingToOriginalType
+        ? originalLlmMapping.fakerMethod
+        : getDefaultFakerMethod(newJsonType);
 
-      const newFakerMethod = newMapping.fakerMethod;
+      if (isSwitchingToOriginalType) {
+        // Restore original LLM mapping
+        onFieldMappingRestored(activeField, originalLlmMapping);
+      } else {
+        // Use default faker method for new type
+        onFieldTypeChanged(activeField, newJsonType);
+        onFieldMethodChanged(activeField, newFakerMethod);
+      }
 
       track('Mock Data JSON Type Changed', {
         field_name: activeField,
@@ -115,17 +120,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,25 +135,20 @@ const FakerSchemaEditorContent = ({
         currentMapping.mongoType === originalLlmMapping.mongoType &&
         newFakerFunction === originalLlmMapping.fakerMethod;
 
-      const newMapping = isSwitchingToLlmSuggestion
-        ? { ...originalLlmMapping }
-        : {
-            ...currentMapping,
-            fakerMethod: newFakerFunction,
-            fakerArgs: [],
-          };
+      if (isSwitchingToLlmSuggestion) {
+        // Restore original LLM mapping
+        onFieldMappingRestored(activeField, originalLlmMapping);
+      } else {
+        // Update the faker method
+        onFieldMethodChanged(activeField, newFakerFunction);
+      }
 
       track('Mock Data Faker Method Changed', {
         field_name: activeField,
         json_type: currentMapping.mongoType,
         previous_faker_method: previousFakerMethod,
         new_faker_method: newFakerFunction,
       });
-
-      setFakerSchemaFormValues({
-        ...fakerSchemaFormValues,
-        [activeField]: newMapping,
-      });
     }
   };
 
@@ -164,7 +159,7 @@ const FakerSchemaEditorContent = ({
           activeField={activeField}
           fields={fieldPaths}
           onFieldSelect={setActiveField}
-          fakerSchema={fakerSchemaFormValues}
+          fakerSchema={fakerSchema}
         />
         {activeJsonType && activeFakerFunction && (
           <FakerMappingSelector
@@ -173,6 +168,11 @@ const FakerSchemaEditorContent = ({
             activeFakerArgs={activeFakerArgs}
             onJsonTypeSelect={onJsonTypeSelect}
             onFakerFunctionSelect={onFakerFunctionSelect}
+            originalLlmFakerMethod={
+              originalLlmResponse[activeField]?.mongoType === activeJsonType
+                ? originalLlmResponse[activeField]?.fakerMethod
+                : undefined
+            }
           />
         )}
       </div>
@@ -191,9 +191,18 @@ const FakerSchemaEditorContent = ({
 const FakerSchemaEditorScreen = ({
   onSchemaConfirmed,
   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 (
     <div data-testid="faker-schema-editor" className={containerStyles}>
@@ -220,11 +229,24 @@ const FakerSchemaEditorScreen = ({
       {fakerSchemaGenerationState.status === 'completed' && (
         <FakerSchemaEditorContent
           fakerSchema={fakerSchemaGenerationState.editedFakerSchema}
+          originalLlmResponse={fakerSchemaGenerationState.originalLlmResponse}
           onSchemaConfirmed={onSchemaConfirmed}
+          onFieldTypeChanged={onFieldTypeChanged}
+          onFieldMethodChanged={onFieldMethodChanged}
+          onFieldMappingRestored={onFieldMappingRestored}
         />
       )}
     </div>
   );
 };
 
-export default connect()(FakerSchemaEditorScreen);
+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(() => ({}), mapDispatchToProps)(FakerSchemaEditorScreen);