feat(compass-collection): Retrigger schema analysis after adding data to an empty collection to enable Mock Data Generator button – CLOUDP-356729

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

@@ -18,6 +18,9 @@ import {
 import { type CollectionMetadata } from 'mongodb-collection-model';
 import type { types } from '@mongodb-js/mdb-experiment-js';
 
+// Wait time in ms for async operations to complete
+const WAIT_TIME = 50;
+
 // Helper function to create proper mock assignment objects for testing
 const createMockAssignment = (
   variant: ExperimentTestGroup
@@ -86,12 +89,41 @@ describe('Collection Tab Content store', function () {
   const sandbox = Sinon.createSandbox();
 
   const localAppRegistry = sandbox.spy(new AppRegistry());
+  const globalAppRegistry = sandbox.spy(new AppRegistry());
   const analyzeCollectionSchemaStub = sandbox
     .stub(collectionTabModule, 'analyzeCollectionSchema')
     .returns(async () => {});
 
-  const dataService = {} as any;
-  const atlasAiService = {} as any;
+  // Track event listeners for proper cleanup
+  const eventListeners: Array<{
+    registry: AppRegistry;
+    eventName: string;
+    listener: (...args: unknown[]) => void;
+  }> = [];
+
+  const mockActivateHelpers = {
+    on: sandbox.spy(
+      (
+        registry: AppRegistry,
+        eventName: string,
+        listener: (...args: unknown[]) => void
+      ) => {
+        registry.on(eventName, listener);
+        eventListeners.push({ registry, eventName, listener });
+      }
+    ),
+    cleanup: sandbox.spy(() => {
+      // Remove all tracked event listeners
+      eventListeners.forEach(({ registry, eventName, listener }) => {
+        registry.removeListener(eventName, listener);
+      });
+      eventListeners.length = 0;
+    }),
+    addCleanup: sandbox.spy(),
+  };
+
+  const dataService = {} as never;
+  const atlasAiService = {} as never;
   let store: ReturnType<typeof activatePlugin>['store'];
   let deactivate: ReturnType<typeof activatePlugin>['deactivate'];
 
@@ -101,7 +133,7 @@ describe('Collection Tab Content store', function () {
     experimentationServices: Partial<ExperimentationServices> = {},
     connectionInfoRef: Partial<
       ReturnType<typeof connectionInfoRefLocator>
-    > = {},
+    > = mockAtlasConnectionInfo,
     logger = createNoopLogger('COMPASS-COLLECTION-TEST'),
     preferences = new ReadOnlyPreferenceAccess({
       enableGenAIFeatures: true,
@@ -128,14 +160,15 @@ describe('Collection Tab Content store', function () {
         dataService,
         atlasAiService,
         localAppRegistry,
-        collection: mockCollection as any,
-        workspaces: workspaces as any,
-        experimentationServices: experimentationServices as any,
-        connectionInfoRef: connectionInfoRef as any,
+        globalAppRegistry,
+        collection: mockCollection as never,
+        workspaces: workspaces as never,
+        experimentationServices: experimentationServices as never,
+        connectionInfoRef: connectionInfoRef as never,
         logger,
         preferences,
       },
-      { on() {}, cleanup() {}, addCleanup() {} } as any
+      mockActivateHelpers as never
     ));
     await waitFor(() => {
       expect(store.getState())
@@ -146,17 +179,22 @@ describe('Collection Tab Content store', function () {
   };
 
   afterEach(function () {
+    mockActivateHelpers.cleanup();
     sandbox.resetHistory();
     deactivate();
   });
 
   describe('selectTab', function () {
     it('should set active tab', async function () {
       const openCollectionWorkspaceSubtab = sandbox.spy();
-      const store = await configureStore(undefined, {
-        openCollectionWorkspaceSubtab,
-      });
-      store.dispatch(selectTab('Documents') as any);
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      const store = await configureStore(
+        undefined,
+        { openCollectionWorkspaceSubtab },
+        { assignExperiment }
+      );
+      store.dispatch(selectTab('Documents') as never);
       expect(openCollectionWorkspaceSubtab).to.have.been.calledWith(
         'workspace-tab-id',
         'Documents'
@@ -206,7 +244,7 @@ describe('Collection Tab Content store', function () {
       );
 
       // Wait a bit to ensure assignment would have happened if it was going to
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(assignExperiment).to.not.have.been.called;
     });
 
@@ -229,7 +267,7 @@ describe('Collection Tab Content store', function () {
       );
 
       // Wait a bit to ensure assignment would have happened if it was going to
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(assignExperiment).to.not.have.been.called;
 
       // Store should still be functional
@@ -278,7 +316,7 @@ describe('Collection Tab Content store', function () {
       );
 
       // Wait a bit to ensure assignment would have happened if it was going to
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(assignExperiment).to.not.have.been.called;
     });
 
@@ -296,7 +334,7 @@ describe('Collection Tab Content store', function () {
       );
 
       // Wait a bit to ensure assignment would have happened if it was going to
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(assignExperiment).to.not.have.been.called;
     });
   });
@@ -377,7 +415,7 @@ describe('Collection Tab Content store', function () {
       });
 
       // Wait a bit to ensure schema analysis would not have been called
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
 
@@ -428,7 +466,7 @@ describe('Collection Tab Content store', function () {
       });
 
       // Wait a bit to ensure schema analysis would not have been called
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
 
@@ -450,7 +488,7 @@ describe('Collection Tab Content store', function () {
       });
 
       // Wait a bit to ensure schema analysis would not have been called
-      await new Promise((resolve) => setTimeout(resolve, 50));
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
       expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
   });
@@ -470,12 +508,195 @@ describe('Collection Tab Content store', function () {
       });
 
       // Dispatch cancel action
-      store.dispatch(collectionTabModule.cancelSchemaAnalysis() as any);
+      store.dispatch(collectionTabModule.cancelSchemaAnalysis() as never);
 
       // Verify the state is reset to initial
-      expect((store.getState() as any).schemaAnalysis.status).to.equal(
-        'initial'
+      expect(
+        (store.getState() as { schemaAnalysis: { status: string } })
+          .schemaAnalysis.status
+      ).to.equal('initial');
+    });
+  });
+
+  describe('document-inserted event listener', function () {
+    it('should re-trigger schema analysis when document is inserted into current collection', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      const store = await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo,
+        undefined,
+        undefined,
+        { ...defaultMetadata, isReadonly: false, isTimeSeries: false }
+      );
+
+      // Wait for initial schema analysis to complete
+      await waitFor(() => {
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
+
+      // Reset the stub to track new calls
+      analyzeCollectionSchemaStub.resetHistory();
+
+      // Simulate the empty collection
+      store.dispatch({
+        type: 'compass-collection/SchemaAnalysisFailed',
+        error: new Error('No documents found'),
+      } as never);
+
+      // Trigger the document-inserted event
+      globalAppRegistry.emit(
+        'document-inserted',
+        {
+          ns: defaultMetadata.namespace,
+          view: 'default',
+          mode: 'default',
+          multiple: false,
+          docs: [{ _id: 'test-doc-id', name: 'test' }],
+        },
+        { connectionId: mockAtlasConnectionInfo.current.id }
+      );
+
+      // Wait for schema analysis to be re-triggered
+      await waitFor(() => {
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
+    });
+
+    it('should not re-trigger schema analysis for different collection', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      // Wait for initial schema analysis to complete
+      await waitFor(() => {
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
+
+      // Reset the stub to track new calls
+      analyzeCollectionSchemaStub.resetHistory();
+
+      // Trigger the document-inserted event with different collection
+      globalAppRegistry.emit(
+        'document-inserted',
+        {
+          ns: 'different.collection',
+          view: 'default',
+          mode: 'default',
+          multiple: false,
+          docs: [{ _id: 'test-doc-id', name: 'test' }],
+        },
+        { connectionId: mockAtlasConnectionInfo.current.id }
       );
+
+      // Wait a bit to ensure schema analysis is not called
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+
+    it('should not re-trigger schema analysis for different connection', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      // Wait for initial schema analysis to complete
+      await waitFor(() => {
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
+
+      // Reset the stub to track new calls
+      analyzeCollectionSchemaStub.resetHistory();
+
+      // Trigger the document-inserted event with different connection
+      globalAppRegistry.emit(
+        'document-inserted',
+        {
+          ns: defaultMetadata.namespace,
+          view: 'default',
+          mode: 'default',
+          multiple: false,
+          docs: [{ _id: 'test-doc-id', name: 'test' }],
+        },
+        { connectionId: 'different-connection-id' }
+      );
+
+      // Wait a bit to ensure schema analysis is not called
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+
+    it('should not re-trigger schema analysis when user is not in experiment variant', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorControl)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      // Wait for initial assignment check
+      await waitFor(() => {
+        expect(getAssignment).to.have.been.calledOnce;
+      });
+
+      // Schema analysis should not have been called initially
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+
+      // Verify the schema analysis state is INITIAL (as expected for control variant)
+      const initialState = store.getState() as {
+        schemaAnalysis: { status: string };
+      };
+      expect(initialState.schemaAnalysis.status).to.equal('initial');
+
+      // Trigger the document-inserted event
+      globalAppRegistry.emit(
+        'document-inserted',
+        {
+          ns: defaultMetadata.namespace,
+          view: 'default',
+          mode: 'default',
+          multiple: false,
+          docs: [{ _id: 'test-doc-id', name: 'test' }],
+        },
+        { connectionId: mockAtlasConnectionInfo.current.id }
+      );
+
+      // Wait a bit to ensure schema analysis is not called
+      await new Promise((resolve) => setTimeout(resolve, WAIT_TIME));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
   });
 });