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

@@ -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,39 +121,52 @@ describe('FakerMappingSelector', () => {
     );
   });
 
-  it('should show warning banner when faker method is unrecognized', () => {
+  it('should always include the original LLM faker method in the dropdown', () => {
+    const originalLlmMethod = 'custom.llmMethod';
+
     render(
       <FakerMappingSelector
-        activeJsonType={mockActiveJsonType}
-        activeFakerFunction={UNRECOGNIZED_FAKER_METHOD}
-        activeFakerArgs={mockActiveFakerArgs}
+        activeJsonType="String"
+        activeFakerFunction="lorem.word"
+        activeFakerArgs={[]}
         onJsonTypeSelect={onJsonTypeSelectStub}
         onFakerFunctionSelect={onFakerFunctionSelectStub}
+        originalLlmFakerMethod={originalLlmMethod}
       />
     );
 
-    expect(
-      screen.getByText(
-        /Please select a function or we will default fill this field/
-      )
-    ).to.exist;
+    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 show warning banner when faker method is recognized', () => {
+  it('should not duplicate the original LLM method if it is already in the standard methods', () => {
+    const originalLlmMethod = 'lorem.word';
+
     render(
       <FakerMappingSelector
-        activeJsonType={mockActiveJsonType}
-        activeFakerFunction={mockActiveFakerFunction}
-        activeFakerArgs={mockActiveFakerArgs}
+        activeJsonType="String"
+        activeFakerFunction="lorem.sentence"
+        activeFakerArgs={[]}
         onJsonTypeSelect={onJsonTypeSelectStub}
         onFakerFunctionSelect={onFakerFunctionSelectStub}
+        originalLlmFakerMethod={originalLlmMethod}
       />
     );
 
-    expect(
-      screen.queryByText(
-        /Please select a function or we will default fill this field/
-      )
-    ).to.not.exist;
+    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

@@ -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,
@@ -67,6 +64,7 @@ interface Props {
   onJsonTypeSelect: (jsonType: MongoDBFieldType) => void;
   activeFakerArgs: FakerArg[];
   onFakerFunctionSelect: (fakerFunction: string) => void;
+  originalLlmFakerMethod?: string;
 }
 
 const FakerMappingSelector = ({
@@ -75,16 +73,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}>
@@ -113,29 +113,17 @@ const FakerMappingSelector = ({
           </Option>
         ))}
       </Select>
-      {activeFakerFunction === UNRECOGNIZED_FAKER_METHOD ? (
-        <Banner variant={BannerVariant.Warning}>
-          Please select a function or we will default fill this field with the
-          string &quot;Unrecognized&quot;
-        </Banner>
-      ) : (
-        <>
-          <Label htmlFor="faker-function-call-preview">
-            Preview Faker Function Call
-          </Label>
-          <Code
-            id="faker-function-call-preview"
-            data-testid="faker-function-call-preview"
-            language="javascript"
-            copyButtonAppearance="none"
-          >
-            {formatFakerFunctionCallWithArgs(
-              activeFakerFunction,
-              activeFakerArgs
-            )}
-          </Code>
-        </>
-      )}
+      <Label htmlFor="faker-function-call-preview">
+        Preview Faker Function Call
+      </Label>
+      <Code
+        id="faker-function-call-preview"
+        data-testid="faker-function-call-preview"
+        language="javascript"
+        copyButtonAppearance="none"
+      >
+        {formatFakerFunctionCallWithArgs(activeFakerFunction, activeFakerArgs)}
+      </Code>
     </div>
   );
 };

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(null, mapDispatchToProps)(FakerSchemaEditorScreen);