feat(compass-collection): Collection Plugin Experimentation Assignment Integration – CLOUDP-333845 (#7165)

jcobis · web-flow · commit db7e9b81432a · 2025-08-07T11:30:54.000-04:00
* WIP

* Org level flag and tests

* WIP

* WIP

* Absolute import

* Remove unneeded check

* Comment address

* Change export

* Comment

* Change exports

* Rename existing test name enum
diff --git a/packages/compass-collection/src/index.ts b/packages/compass-collection/src/index.ts
@@ -6,10 +6,14 @@ import {
   dataServiceLocator,
   type DataServiceLocator,
   type DataService,
+  connectionInfoRefLocator,
 } from '@mongodb-js/compass-connections/provider';
 import { collectionModelLocator } from '@mongodb-js/compass-app-stores/provider';
 import type { WorkspacePlugin } from '@mongodb-js/compass-workspaces';
 import { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
+import { preferencesLocator } from 'compass-preferences-model/provider';
 import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
@@ -29,6 +33,10 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
       dataService: dataServiceLocator as DataServiceLocator<keyof DataService>,
       collection: collectionModelLocator,
       workspaces: workspacesServiceLocator,
+      experimentationServices: experimentationServiceLocator,
+      connectionInfoRef: connectionInfoRefLocator,
+      logger: createLoggerLocator('COMPASS-COLLECTION'),
+      preferences: preferencesLocator,
     }
   ),
   content: CollectionTab,
diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts
@@ -5,6 +5,7 @@ import type AppRegistry from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
 import type { CollectionSubtab } from '@mongodb-js/compass-workspaces';
 import type { DataService } from '@mongodb-js/compass-connections/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -20,6 +21,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     localAppRegistry: AppRegistry;
     dataService: DataService;
     workspaces: ReturnType<typeof workspacesServiceLocator>;
+    experimentationServices: ReturnType<typeof experimentationServiceLocator>;
   },
   A
 >;
diff --git a/packages/compass-collection/src/stores/collection-tab.spec.ts b/packages/compass-collection/src/stores/collection-tab.spec.ts
@@ -6,6 +6,11 @@ import Sinon from 'sinon';
 import AppRegistry from '@mongodb-js/compass-app-registry';
 import { expect } from 'chai';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry';
+import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
+import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
+import { ReadOnlyPreferenceAccess } from 'compass-preferences-model/provider';
+import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
 
 const defaultMetadata = {
   namespace: 'test.foo',
@@ -32,6 +37,32 @@ const mockCollection = {
   },
 };
 
+const mockAtlasConnectionInfo = {
+  current: {
+    id: 'test-connection',
+    title: 'Test Connection',
+    connectionOptions: {
+      connectionString: 'mongodb://localhost:27017',
+    },
+    atlasMetadata: {
+      clusterName: 'test-cluster',
+      projectId: 'test-project',
+      orgId: 'test-org',
+      clusterUniqueId: 'test-cluster-unique-id',
+      clusterType: 'REPLICASET' as const,
+      clusterState: 'IDLE' as const,
+      metricsId: 'test-metrics-id',
+      metricsType: 'replicaSet' as const,
+      regionalBaseUrl: null,
+      instanceSize: 'M10',
+      supports: {
+        globalWrites: false,
+        rollingIndexes: true,
+      },
+    },
+  },
+};
+
 describe('Collection Tab Content store', function () {
   const sandbox = Sinon.createSandbox();
 
@@ -42,7 +73,19 @@ describe('Collection Tab Content store', function () {
 
   const configureStore = async (
     options: Partial<CollectionTabOptions> = {},
-    workspaces: Partial<ReturnType<typeof workspacesServiceLocator>> = {}
+    workspaces: Partial<ReturnType<typeof workspacesServiceLocator>> = {},
+    experimentationServices: Partial<
+      ReturnType<typeof experimentationServiceLocator>
+    > = {},
+    connectionInfoRef: Partial<
+      ReturnType<typeof connectionInfoRefLocator>
+    > = {},
+    logger = createNoopLogger('COMPASS-COLLECTION-TEST'),
+    preferences = new ReadOnlyPreferenceAccess({
+      enableGenAIFeatures: true,
+      enableGenAIFeaturesAtlasOrg: true,
+      cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
+    })
   ) => {
     ({ store, deactivate } = activatePlugin(
       {
@@ -54,6 +97,10 @@ describe('Collection Tab Content store', function () {
         localAppRegistry,
         collection: mockCollection as any,
         workspaces: workspaces as any,
+        experimentationServices: experimentationServices as any,
+        connectionInfoRef: connectionInfoRef as any,
+        logger,
+        preferences,
       },
       { on() {}, cleanup() {} } as any
     ));
@@ -76,11 +123,112 @@ describe('Collection Tab Content store', function () {
       const store = await configureStore(undefined, {
         openCollectionWorkspaceSubtab,
       });
-      store.dispatch(selectTab('Documents'));
+      store.dispatch(selectTab('Documents') as any);
       expect(openCollectionWorkspaceSubtab).to.have.been.calledWith(
         'workspace-tab-id',
         'Documents'
       );
     });
   });
+
+  describe('experimentation integration', function () {
+    it('should assign experiment when Atlas metadata is available', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(assignExperiment).to.have.been.calledOnceWith(
+          ExperimentTestName.mockDataGenerator,
+          {
+            team: 'Atlas Growth',
+          }
+        );
+      });
+    });
+
+    it('should not assign experiment when Atlas metadata is missing', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+      const mockConnectionInfoRef = {
+        current: {
+          id: 'test-connection',
+          title: 'Test Connection',
+          connectionOptions: {
+            connectionString: 'mongodb://localhost:27017',
+          },
+          // No atlasMetadata
+        },
+      };
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockConnectionInfoRef
+      );
+
+      // Wait a bit to ensure assignment would have happened if it was going to
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(assignExperiment).to.not.have.been.called;
+    });
+
+    it('should not assign experiment when AI features are disabled at the org level', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      const mockPreferences = new ReadOnlyPreferenceAccess({
+        enableGenAIFeatures: true,
+        enableGenAIFeaturesAtlasOrg: false, // Disabled at org level
+        cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
+      });
+
+      const store = await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo,
+        undefined,
+        mockPreferences
+      );
+
+      // Wait a bit to ensure assignment would have happened if it was going to
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(assignExperiment).to.not.have.been.called;
+
+      // Store should still be functional
+      await waitFor(() => {
+        expect(store.getState())
+          .to.have.property('metadata')
+          .deep.eq(defaultMetadata);
+      });
+    });
+
+    it('should handle assignment errors gracefully', async function () {
+      const assignExperiment = sandbox.spy(() =>
+        Promise.reject(new Error('Assignment failed'))
+      );
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(assignExperiment).to.have.been.calledOnce;
+      });
+
+      // Store should still be functional despite assignment error
+      await waitFor(() => {
+        expect(store.getState())
+          .to.have.property('metadata')
+          .deep.eq(defaultMetadata);
+      });
+    });
+  });
 });
diff --git a/packages/compass-collection/src/stores/collection-tab.ts b/packages/compass-collection/src/stores/collection-tab.ts
@@ -9,6 +9,14 @@ import reducer, {
 import type { Collection } from '@mongodb-js/compass-app-stores/provider';
 import type { ActivateHelpers } from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
+import type { Logger } from '@mongodb-js/compass-logging/provider';
+import {
+  isAIFeatureEnabled,
+  type PreferencesAccess,
+} from 'compass-preferences-model/provider';
+import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
 
 export type CollectionTabOptions = {
   /**
@@ -31,18 +39,29 @@ export type CollectionTabServices = {
   collection: Collection;
   localAppRegistry: AppRegistry;
   workspaces: ReturnType<typeof workspacesServiceLocator>;
+  experimentationServices: ReturnType<typeof experimentationServiceLocator>;
+  connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
+  logger: Logger;
+  preferences: PreferencesAccess;
 };
 
 export function activatePlugin(
   { namespace, editViewName, tabId }: CollectionTabOptions,
   services: CollectionTabServices,
   { on, cleanup }: ActivateHelpers
-) {
+): {
+  store: ReturnType<typeof createStore>;
+  deactivate: () => void;
+} {
   const {
     dataService,
     collection: collectionModel,
     localAppRegistry,
     workspaces,
+    experimentationServices,
+    connectionInfoRef,
+    logger,
+    preferences,
   } = services;
 
   if (!collectionModel) {
@@ -64,6 +83,7 @@ export function activatePlugin(
         dataService,
         workspaces,
         localAppRegistry,
+        experimentationServices,
       })
     )
   );
@@ -86,6 +106,25 @@ export function activatePlugin(
 
   void collectionModel.fetchMetadata({ dataService }).then((metadata) => {
     store.dispatch(collectionMetadataFetched(metadata));
+
+    // Assign experiment for Mock Data Generator
+    // Only assign when we're connected to Atlas and the org-level setting for AI features is enabled
+    if (
+      connectionInfoRef.current?.atlasMetadata?.clusterName && // Ensures we only assign in Atlas
+      isAIFeatureEnabled(preferences.getPreferences()) // Ensures org-level AI features setting is enabled
+    ) {
+      void experimentationServices
+        .assignExperiment(ExperimentTestName.mockDataGenerator, {
+          team: 'Atlas Growth',
+        })
+        .catch((error) => {
+          logger.debug('Mock Data Generator experiment assignment failed', {
+            experiment: ExperimentTestName.mockDataGenerator,
+            namespace: namespace,
+            error: error instanceof Error ? error.message : String(error),
+          });
+        });
+    }
   });
 
   return {
diff --git a/packages/compass-indexes/src/components/create-index-modal/create-index-modal.tsx b/packages/compass-indexes/src/components/create-index-modal/create-index-modal.tsx
@@ -24,7 +24,7 @@ import {
   useTrackOnChange,
   type TrackFunction,
   useFireExperimentViewed,
-  TestName,
+  ExperimentTestName,
   useTelemetry,
 } from '@mongodb-js/compass-telemetry/provider';
 import { useConnectionInfoRef } from '@mongodb-js/compass-connections/provider';
@@ -91,7 +91,7 @@ function CreateIndexModal({
     usePreference('showIndexesGuidanceVariant') && enableInIndexesGuidanceExp;
 
   useFireExperimentViewed({
-    testName: TestName.earlyJourneyIndexesGuidance,
+    testName: ExperimentTestName.earlyJourneyIndexesGuidance,
     shouldFire: enableInIndexesGuidanceExp && isVisible,
   });
 
diff --git a/packages/compass-telemetry/src/experimentation-provider.tsx b/packages/compass-telemetry/src/experimentation-provider.tsx
@@ -1,15 +1,16 @@
 import React, { createContext, useContext, useRef } from 'react';
 import type { types } from '@mongodb-js/mdb-experiment-js';
 import type { typesReact } from '@mongodb-js/mdb-experiment-js/react';
+import type { ExperimentTestName } from './growth-experiments';
 
 type UseAssignmentHook = (
-  experimentName: string,
+  experimentName: ExperimentTestName,
   trackIsInSample: boolean,
   options?: typesReact.UseAssignmentOptions<types.TypeData>
 ) => typesReact.UseAssignmentResponse<types.TypeData>;
 
 type AssignExperimentFn = (
-  experimentName: string,
+  experimentName: ExperimentTestName,
   options?: types.AssignOptions<string>
 ) => Promise<types.AsyncStatus | null>;
 
@@ -34,7 +35,7 @@ const initialContext: CompassExperimentationProviderContextValue = {
   },
 };
 
-const ExperimentationContext =
+export const ExperimentationContext =
   createContext<CompassExperimentationProviderContextValue>(initialContext);
 
 // Provider component that accepts MMS experiment utils as props
diff --git a/packages/compass-telemetry/src/growth-experiments.ts b/packages/compass-telemetry/src/growth-experiments.ts
@@ -1,3 +1,4 @@
-export enum TestName {
+export enum ExperimentTestName {
   earlyJourneyIndexesGuidance = 'EARLY_JOURNEY_INDEXES_GUIDANCE_20250328',
+  mockDataGenerator = 'MOCK_DATA_GENERATOR_20251001',
 }
diff --git a/packages/compass-telemetry/src/index.ts b/packages/compass-telemetry/src/index.ts
@@ -7,3 +7,5 @@ export type {
 } from './types';
 
 export { CompassExperimentationProvider } from './experimentation-provider';
+export { experimentationServiceLocator } from './provider';
+export { ExperimentTestName } from './growth-experiments';
diff --git a/packages/compass-telemetry/src/provider.tsx b/packages/compass-telemetry/src/provider.tsx

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
`1`		`-export enum TestName {`
	`1`	`+export enum ExperimentTestName {`
`2`	`2`	`earlyJourneyIndexesGuidance = 'EARLY_JOURNEY_INDEXES_GUIDANCE_20250328',`
	`3`	`+ mockDataGenerator = 'MOCK_DATA_GENERATOR_20251001',`
`3`	`4`	`}`