feat(compass-collection) Disclaimer Screen - Mock Data Generator CLOUDP-333851

Author: ncarbon

package-lock.json

@@ -52,6 +52,7 @@
     "@mongodb-js/compass-app-stores": "^7.56.0",
     "@mongodb-js/compass-components": "^1.48.0",
     "@mongodb-js/compass-connections": "^1.70.0",
+    "@mongodb-js/compass-generative-ai": "^0.50.0",
     "@mongodb-js/compass-logging": "^1.7.11",
     "@mongodb-js/compass-telemetry": "^1.13.0",
     "@mongodb-js/compass-workspaces": "^0.51.0",

packages/compass-collection/package.json

@@ -22,17 +22,24 @@ import CollectionHeaderActions from '../collection-header-actions';
 describe('CollectionHeaderActions [Component]', function () {
   let preferences: PreferencesAccess;
   let mockUseAssignment: sinon.SinonStub;
+  let mockGetAssignment: sinon.SinonStub;
 
   beforeEach(async function () {
     preferences = await createSandboxFromDefaultPreferences();
     mockUseAssignment = sinon.stub();
+    mockGetAssignment = sinon.stub();
     mockUseAssignment.returns({
       assignment: {
         assignmentData: {
           variant: 'mockDataGeneratorControl',
         },
       },
     });
+    mockGetAssignment.returns({
+      assignmentData: {
+        variant: 'mockDataGeneratorControl',
+      },
+    });
   });
 
   afterEach(function () {
@@ -48,6 +55,7 @@ describe('CollectionHeaderActions [Component]', function () {
       <CompassExperimentationProvider
         useAssignment={mockUseAssignment}
         assignExperiment={sinon.stub()}
+        getAssignment={mockGetAssignment}
       >
         <WorkspacesServiceProvider
           value={workspaceService as WorkspacesService}

packages/compass-collection/src/components/collection-header-actions/collection-header-actions.spec.tsx

@@ -23,7 +23,7 @@ describe('MockDataGeneratorModal', () => {
 
   function renderModal({
     isOpen = true,
-    currentStep = MockDataGeneratorStep.AI_DISCLAIMER,
+    currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION,
   } = {}) {
     return render(
       <MockDataGeneratorModal

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

@@ -70,7 +70,7 @@ const MockDataGeneratorModal = ({
       <ModalFooter className={footerStyles}>
         <Button
           onClick={onPreviousStep}
-          disabled={currentStep === MockDataGeneratorStep.AI_DISCLAIMER}
+          disabled={currentStep === MockDataGeneratorStep.SCHEMA_CONFIRMATION}
         >
           Back
         </Button>

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

@@ -18,6 +18,7 @@ import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
 } from './plugin-tab-title';
+import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
 
 export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
   name: CollectionWorkspaceTitle,
@@ -37,6 +38,7 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
       connectionInfoRef: connectionInfoRefLocator,
       logger: createLoggerLocator('COMPASS-COLLECTION'),
       preferences: preferencesLocator,
+      atlasAiService: atlasAiServiceLocator,
     }
   ),
   content: CollectionTab,

packages/compass-collection/src/index.ts

@@ -10,6 +10,8 @@ import type { DataService } from '@mongodb-js/compass-connections/provider';
 import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
 import { type PreferencesAccess } from 'compass-preferences-model/provider';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
+
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -22,11 +24,11 @@ import {
 } from '../schema-analysis-types';
 import { calculateSchemaDepth } from '../calculate-schema-depth';
 import type { Document, MongoError } from 'mongodb';
+import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 const DEFAULT_SAMPLE_SIZE = 100;
 
 const NO_DOCUMENTS_ERROR = 'No documents found in the collection to analyze.';
-import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -63,6 +65,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
     logger: Logger;
     preferences: PreferencesAccess;
+    atlasAiService: AtlasAiService;
   },
   A
 >;
@@ -89,6 +92,7 @@ enum CollectionActions {
   MockDataGeneratorModalClosed = 'compass-collection/MockDataGeneratorModalClosed',
   MockDataGeneratorNextButtonClicked = 'compass-collection/MockDataGeneratorNextButtonClicked',
   MockDataGeneratorPreviousButtonClicked = 'compass-collection/MockDataGeneratorPreviousButtonClicked',
+  AiDisclaimerModalShown = 'compass-collection/AiDisclaimerModalShown',
 }
 
 interface CollectionMetadataFetchedAction {
@@ -104,6 +108,10 @@ interface SchemaAnalysisStartedAction {
   type: CollectionActions.SchemaAnalysisStarted;
 }
 
+interface AiDisclaimerModalShownAction {
+  type: CollectionActions.AiDisclaimerModalShown;
+}
+
 interface SchemaAnalysisFinishedAction {
   type: CollectionActions.SchemaAnalysisFinished;
   schema: Schema;
@@ -146,7 +154,7 @@ const reducer: Reducer<CollectionState, Action> = (
     },
     mockDataGenerator: {
       isModalOpen: false,
-      currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+      currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     },
   },
   action
@@ -234,7 +242,7 @@ const reducer: Reducer<CollectionState, Action> = (
       mockDataGenerator: {
         ...state.mockDataGenerator,
         isModalOpen: true,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     };
   }
@@ -359,6 +367,10 @@ export const mockDataGeneratorPreviousButtonClicked =
     return { type: CollectionActions.MockDataGeneratorPreviousButtonClicked };
   };
 
+export const aIDisclaimerModalOpened = (): AiDisclaimerModalShownAction => {
+  return { type: CollectionActions.AiDisclaimerModalShown };
+};
+
 export const selectTab = (
   tabName: CollectionSubtab
 ): CollectionThunkAction<void> => {
@@ -370,6 +382,32 @@ export const selectTab = (
   };
 };
 
+export const openMockDataGeneratorModal = (): CollectionThunkAction<
+  Promise<void>
+> => {
+  return async (
+    dispatch,
+    _getState,
+    { atlasAiService, preferences, logger }
+  ) => {
+    try {
+      if (preferences.getPreferences().optInGenAIFeatures) {
+        // If the user has already opted in, we can directly open the modal
+        dispatch(mockDataGeneratorModalOpened());
+      } else if (process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN !== 'true') {
+        await atlasAiService.ensureAiFeatureAccess();
+      }
+    } catch (error) {
+      logger.log.error(
+        mongoLogId(1_001_000_364),
+        'Collections',
+        'Failed to ensure AI feature access',
+        error
+      );
+    }
+  };
+};
+
 export const analyzeCollectionSchema = (): CollectionThunkAction<
   Promise<void>
 > => {

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

@@ -63,6 +63,7 @@ describe('Collection Tab Content store', function () {
     .stub(collectionTabModule, 'analyzeCollectionSchema')
     .returns(async () => {});
   const dataService = {} as any;
+  const atlasAiService = {} as any;
   let store: ReturnType<typeof activatePlugin>['store'];
   let deactivate: ReturnType<typeof activatePlugin>['deactivate'];
 
@@ -106,6 +107,7 @@ describe('Collection Tab Content store', function () {
         connectionInfoRef: connectionInfoRef as any,
         logger,
         preferences,
+        atlasAiService,
       },
       { on() {}, cleanup() {} } as any
     ));

packages/compass-collection/src/stores/collection-tab.spec.ts

@@ -7,6 +7,7 @@ import reducer, {
   selectTab,
   collectionMetadataFetched,
   analyzeCollectionSchema,
+  openMockDataGeneratorModal,
 } from '../modules/collection-tab';
 import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
@@ -20,8 +21,12 @@ import {
   isAIFeatureEnabled,
   type PreferencesAccess,
 } from 'compass-preferences-model/provider';
-import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
+import {
+  ExperimentTestGroup,
+  ExperimentTestName,
+} from '@mongodb-js/compass-telemetry/provider';
 import { SCHEMA_ANALYSIS_STATE_INITIAL } from '../schema-analysis-types';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
 
 export type CollectionTabOptions = {
   /**
@@ -48,6 +53,7 @@ export type CollectionTabServices = {
   connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
   logger: Logger;
   preferences: PreferencesAccess;
+  atlasAiService: AtlasAiService;
 };
 
 export function activatePlugin(
@@ -67,6 +73,7 @@ export function activatePlugin(
     connectionInfoRef,
     logger,
     preferences,
+    atlasAiService,
   } = services;
 
   if (!collectionModel) {
@@ -87,7 +94,7 @@ export function activatePlugin(
       },
       mockDataGenerator: {
         isModalOpen: false,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     },
     applyMiddleware(
@@ -98,6 +105,7 @@ export function activatePlugin(
         experimentationServices,
         logger,
         preferences,
+        atlasAiService,
       })
     )
   );
@@ -118,6 +126,28 @@ export function activatePlugin(
     store.dispatch(selectTab('Schema'));
   });
 
+  on(localAppRegistry, 'open-mock-data-generator-modal', () => {
+    void experimentationServices
+      .getAssignment(ExperimentTestName.mockDataGenerator, {
+        team: 'Atlas Growth',
+      })
+      .then((mockDataGeneratorAssignment) => {
+        if (
+          mockDataGeneratorAssignment?.assignmentData?.variant ===
+          ExperimentTestGroup.mockDataGeneratorVariant
+        ) {
+          void store.dispatch(openMockDataGeneratorModal());
+        }
+      })
+      .catch((error) => {
+        logger.debug('Mock Data Generator get experiment assignment', {
+          experiment: ExperimentTestName.mockDataGenerator,
+          namespace: namespace,
+          error: error instanceof Error ? error.message : String(error),
+        });
+      });
+  });
+
   void collectionModel.fetchMetadata({ dataService }).then((metadata) => {
     store.dispatch(collectionMetadataFetched(metadata));