clean up probability handling in AtlasAiService and in faker schema validation

Author: ncarbon

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx

@@ -5,6 +5,7 @@ import {
   renderWithActiveConnection,
   waitFor,
   userEvent,
+  waitForElementToBeRemoved,
 } from '@mongodb-js/testing-library-compass';
 import { Provider } from 'react-redux';
 import { createStore, applyMiddleware } from 'redux';
@@ -511,10 +512,9 @@ describe('MockDataGeneratorModal', () => {
 
       // advance to the schema editor step
       userEvent.click(screen.getByText('Confirm'));
-
-      await waitFor(() => {
-        expect(screen.getByTestId('faker-schema-editor')).to.exist;
-      });
+      await waitForElementToBeRemoved(() =>
+        screen.queryByTestId('faker-schema-editor-loader')
+      );
 
       // select the "name" field
       userEvent.click(screen.getByText('name'));

packages/compass-collection/src/modules/collection-tab.ts

@@ -772,9 +772,12 @@ const validateFakerSchema = (
 
   // Process all input schema fields in a single O(n) pass
   for (const fieldPath of Object.keys(inputSchema)) {
-    const fakerMapping = fakerSchemaRaw[fieldPath];
-
-    if (fakerMapping) {
+    if (fakerSchemaRaw[fieldPath]) {
+      // input schema field exists in faker schema
+      const fakerMapping = {
+        ...fakerSchemaRaw[fieldPath],
+        probability: inputSchema[fieldPath].probability,
+      };
       // Validate the faker method
       if (isValidFakerMethod(fakerMapping.fakerMethod)) {
         result[fieldPath] = fakerMapping;
@@ -789,7 +792,7 @@ const validateFakerSchema = (
           mongoType: fakerMapping.mongoType,
           fakerMethod: UNRECOGNIZED_FAKER_METHOD,
           fakerArgs: [],
-          probability: inputSchema[fieldPath].probability,
+          probability: fakerMapping.probability,
         };
       }
     } else {

packages/compass-generative-ai/src/atlas-ai-service.spec.ts

@@ -393,12 +393,10 @@ describe('AtlasAiService', function () {
             name: {
               type: 'string',
               sampleValues: ['John', 'Jane', 'Bob'],
-              probability: 0.9,
             },
             age: {
               type: 'number',
               sampleValues: [25, 30, 35],
-              probability: 0.8,
             },
           },
           includeSampleValues: false,
@@ -536,7 +534,6 @@ describe('AtlasAiService', function () {
             );
             expect(requestBody.schema.age).to.not.have.property('sampleValues');
             expect(requestBody.schema.name.type).to.equal('string');
-            expect(requestBody.schema.age.probability).to.equal(0.8);
           });
 
           it('makes POST request with correct headers and body structure', async function () {

packages/compass-generative-ai/src/atlas-ai-service.ts

@@ -218,7 +218,6 @@ const aiURLConfig = {
 export interface MockDataSchemaRawField {
   type: string;
   sampleValues?: unknown[];
-  probability?: number;
 }
 
 export interface MockDataSchemaRequest {
@@ -475,7 +474,7 @@ export class AtlasAiService {
         Omit<MockDataSchemaRawField, 'sampleValues'>
       > = {};
       for (const [k, v] of Object.entries(schema)) {
-        newSchema[k] = { type: v.type, probability: v.probability };
+        newSchema[k] = { type: v.type };
       }
       schema = newSchema;
     }