chore(e2e-tests): add gen ai test with logged in atlas user COMPASS-8395

packages/compass-e2e-tests/helpers/commands/index.ts

@@ -62,6 +62,7 @@ export * from './unhide-index';
 export * from './hide-visible-modal';
 export * from './hide-visible-toasts';
 export * from './sidebar-collection';
+export * from './switch-pipeline-mode';
 export * from './read-first-document-content';
 export * from './read-stage-operators';
 export * from './click-confirmation-action';

packages/compass-e2e-tests/helpers/commands/set-feature.ts

@@ -1,11 +1,34 @@
 import type { CompassBrowser } from '../compass-browser';
-import type { UserPreferences } from 'compass-preferences-model';
+import type {
+  AllPreferences,
+  UserPreferences,
+} from 'compass-preferences-model';
+import { isTestingWeb } from '../test-runner-context';
 
 export async function setFeature<K extends keyof UserPreferences>(
   browser: CompassBrowser,
   name: K,
   value: UserPreferences[K]
 ): Promise<void> {
+  if (isTestingWeb()) {
+    // When running in Compass web we cannot use the IPC to update the
+    // preferences so we use a global function.
+    await browser.execute(
+      async (_name, _value) => {
+        const attributes: Partial<AllPreferences> = {
+          [_name]: _value === null ? undefined : _value,
+        };
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        await (globalThis as any).__compassWebE2ETestSavePreferences(
+          attributes
+        );
+      },
+      name,
+      value
+    );
+    return;
+  }
+
   await browser.execute(
     async (_name, _value) => {
       // eslint-disable-next-line @typescript-eslint/no-var-requires

packages/compass-e2e-tests/helpers/commands/switch-pipeline-mode.ts

@@ -0,0 +1,10 @@
+import type { CompassBrowser } from '../compass-browser';
+import * as Selectors from '../selectors';
+
+export async function switchPipelineMode(
+  browser: CompassBrowser,
+  mode: 'as-text' | 'builder-ui'
+): Promise<void> {
+  await browser.clickVisible(Selectors.aggregationPipelineModeToggle(mode));
+  await browser.waitForAnimations(Selectors.AggregationBuilderWorkspace);
+}

packages/compass-e2e-tests/helpers/selectors.ts

@@ -1239,12 +1239,11 @@ export const queryBarExportToLanguageButton = (tabName: string): string => {
   const tabSelector = collectionContent(tabName);
   return `${tabSelector} [data-testid="query-bar-open-export-to-language-button"]`;
 };
-export const QueryBarAIEntryButton =
-  '[data-testid="open-ai-query-entry-button"]';
-export const QueryBarAITextInput = '[data-testid="ai-user-text-input"]';
-export const QueryBarAIGenerateQueryButton =
-  '[data-testid="ai-generate-button"]';
-export const QueryBarAIErrorMessageBanner = '[data-testid="ai-error-msg"]';
+export const GenAIEntryButton = '[data-testid="open-ai-query-entry-button"]';
+export const GenAITextInput = '[data-testid="ai-user-text-input"]';
+export const GenAIGenerateQueryButton = '[data-testid="ai-generate-button"]';
+export const GenAIErrorMessageBanner = '[data-testid="ai-error-msg"]';
+export const GenAIOpenButton = '[data-testid="open-gen-ai-button"]';
 
 // Workspace tabs
 export const WorkspaceTabsContainer =

packages/compass-e2e-tests/helpers/test-runner-context.ts

@@ -344,6 +344,10 @@ process.env.HADRON_DISTRIBUTION ??= context.hadronDistribution;
 process.env.COMPASS_WEB_HTTP_PROXY_CLOUD_CONFIG ??=
   context.atlasCloudSandboxCloudConfig ?? 'dev';
 
+if (isTestingAtlasCloudSandbox(context)) {
+  process.env.E2E_TEST_CLOUD_WEB_ENABLE_PREFERENCE_SAVING ??= 'true';
+}
+
 const testServerVersion =
   process.env.MONGODB_VERSION ?? process.env.MONGODB_RUNNER_VERSION;
 

packages/compass-e2e-tests/tests/atlas-cloud/collection-ai-query.test.ts

@@ -0,0 +1,188 @@
+import { expect } from 'chai';
+
+import type { CompassBrowser } from '../../helpers/compass-browser';
+import {
+  init,
+  cleanup,
+  screenshotIfFailed,
+  DEFAULT_CONNECTION_NAME_1,
+} from '../../helpers/compass';
+import type { Compass } from '../../helpers/compass';
+import * as Selectors from '../../helpers/selectors';
+import { createNumbersCollection } from '../../helpers/insert-data';
+import { isTestingAtlasCloudSandbox } from '../../helpers/test-runner-context';
+import { switchPipelineMode } from '../../helpers/commands/switch-pipeline-mode';
+
+describe('Collection ai query', function () {
+  let compass: Compass;
+  let browser: CompassBrowser;
+
+  before(function () {
+    if (!isTestingAtlasCloudSandbox()) {
+      this.skip();
+    }
+  });
+
+  afterEach(async function () {
+    await screenshotIfFailed(compass, this.currentTest);
+    await cleanup(compass);
+  });
+
+  describe('when the feature is enabled', function () {
+    beforeEach(async function () {
+      compass = await init(this.test?.fullTitle());
+      browser = compass.browser;
+      await browser.setupDefaultConnections();
+
+      await createNumbersCollection();
+      await browser.connectToDefaults();
+      await browser.navigateToCollectionTab(
+        DEFAULT_CONNECTION_NAME_1,
+        'test',
+        'numbers',
+        'Documents'
+      );
+
+      await browser.setFeature('enableGenAIFeaturesAtlasProject', true);
+      await browser.setFeature(
+        'enableGenAISampleDocumentPassingOnAtlasProject',
+        true
+      );
+      await browser.setFeature('enableGenAIFeaturesAtlasOrg', true);
+      await browser.setFeature('optInDataExplorerGenAIFeatures', true);
+    });
+
+    describe('on the documents tab', function () {
+      beforeEach(async function () {
+        await browser.navigateToCollectionTab(
+          DEFAULT_CONNECTION_NAME_1,
+          'test',
+          'numbers',
+          'Documents'
+        );
+      });
+
+      it('should update the query bar with a generated query', async function () {
+        // Click the ai entry button.
+        await browser.clickVisible(Selectors.GenAIEntryButton);
+
+        // Enter the ai prompt.
+        await browser.clickVisible(Selectors.GenAITextInput);
+
+        const testUserInput = 'find all documents where i is greater than 50';
+        await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
+
+        // Click generate.
+        await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
+
+        // Wait for the ipc events to succeed.
+        await browser.waitUntil(async function () {
+          // Make sure the query bar was updated.
+          const queryBarFilterContent = await browser.getCodemirrorEditorText(
+            Selectors.queryBarOptionInputFilter('Documents')
+          );
+          return (
+            queryBarFilterContent.includes('$gt') &&
+            queryBarFilterContent.includes('50')
+          );
+        });
+
+        // Run it and check that the correct documents are shown.
+        await browser.runFind('Documents', true);
+        const modifiedResult = await browser.getFirstListDocument();
+        expect(modifiedResult.i).to.be.equal('51');
+      });
+    });
+
+    describe('on the aggregations tab', function () {
+      beforeEach(async function () {
+        await browser.navigateToCollectionTab(
+          DEFAULT_CONNECTION_NAME_1,
+          'test',
+          'numbers',
+          'Aggregations'
+        );
+
+        await switchPipelineMode(browser, 'as-text');
+      });
+
+      it('should update the aggregation editor with a generated aggregation', async function () {
+        // Click the ai entry button.
+        await browser.clickVisible(Selectors.GenAIOpenButton);
+
+        // Enter the ai prompt.
+        await browser.clickVisible(Selectors.GenAITextInput);
+
+        const testUserInput = 'find all documents where i is 99';
+        await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
+
+        // Click generate.
+        await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
+
+        await browser.waitUntil(async function () {
+          const textEditor = browser.$(Selectors.AggregationAsTextEditor);
+          const textContent = await textEditor.getText();
+          return textContent.includes('$match');
+        });
+
+        // Run it and check that the correct documents are shown.
+        await browser.clickVisible(Selectors.RunPipelineButton);
+        const resultsWorkspace = browser.$(
+          Selectors.AggregationResultsWorkspace
+        );
+        await resultsWorkspace.waitForDisplayed();
+
+        await browser.clickVisible(
+          Selectors.AggregationResultsJSONListSwitchButton
+        );
+        const rawDocuments = await browser.getCodemirrorEditorTextAll(
+          Selectors.DocumentJSONEntry
+        );
+        const documents = rawDocuments.map((text) => {
+          return JSON.parse(text);
+        });
+
+        expect(documents).to.have.lengthOf(1);
+        expect(documents[0]).to.have.property('_id');
+        expect(documents[0]).to.have.property('i', 99);
+      });
+    });
+  });
+
+  describe('when the org feature is disabled', function () {
+    beforeEach(async function () {
+      compass = await init(this.test?.fullTitle());
+      browser = compass.browser;
+      await browser.setupDefaultConnections();
+
+      await createNumbersCollection();
+      await browser.connectToDefaults();
+      await browser.navigateToCollectionTab(
+        DEFAULT_CONNECTION_NAME_1,
+        'test',
+        'numbers',
+        'Documents'
+      );
+
+      await browser.setFeature('enableGenAIFeaturesAtlasProject', true);
+      await browser.setFeature(
+        'enableGenAISampleDocumentPassingOnAtlasProject',
+        true
+      );
+      await browser.setFeature('enableGenAIFeaturesAtlasOrg', false);
+      await browser.setFeature('optInDataExplorerGenAIFeatures', true);
+    });
+
+    it('should not show the gen ai intro button', async function () {
+      // Ensure the query bar is shown.
+      await browser
+        .$(Selectors.queryBarOptionInputFilter('Documents'))
+        .waitForDisplayed();
+
+      const aiIntroButton = browser.$(Selectors.GenAIEntryButton);
+      const isSidebarCreateCollectionButtonExisting =
+        await aiIntroButton.isExisting();
+      expect(isSidebarCreateCollectionButtonExisting).to.be.equal(false);
+    });
+  });
+});

packages/compass-e2e-tests/tests/atlas-login.test.ts

@@ -295,7 +295,7 @@ describe('Atlas Login', function () {
       // control over it
       await browser.clickVisible('span=Not now');
 
-      const aiInput = browser.$(Selectors.QueryBarAITextInput);
+      const aiInput = browser.$(Selectors.GenAITextInput);
       expect(await aiInput.isExisting()).to.eq(false);
       expect(await generateQueryButton.isDisplayed()).to.eq(true);
     });
@@ -328,7 +328,7 @@ describe('Atlas Login', function () {
       // control over it
       await browser.clickVisible('span=Not now');
 
-      const aiInput = browser.$(Selectors.QueryBarAITextInput);
+      const aiInput = browser.$(Selectors.GenAITextInput);
       expect(await aiInput.isExisting()).to.eq(false);
       expect(await generateQueryButton.isDisplayed()).to.eq(true);
     });

packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts

@@ -19,6 +19,7 @@ import {
 import { saveAggregationPipeline } from '../helpers/commands/save-aggregation-pipeline';
 import { Key } from 'webdriverio';
 import type { ChainablePromiseElement } from 'webdriverio';
+import { switchPipelineMode } from '../helpers/commands/switch-pipeline-mode';
 
 const { expect } = chai;
 
@@ -83,14 +84,6 @@ async function waitForTab(browser: CompassBrowser, namespace: string) {
   );
 }
 
-async function switchPipelineMode(
-  browser: CompassBrowser,
-  mode: 'as-text' | 'builder-ui'
-) {
-  await browser.clickVisible(Selectors.aggregationPipelineModeToggle(mode));
-  await browser.waitForAnimations(Selectors.AggregationBuilderWorkspace);
-}
-
 async function deleteStage(
   browser: CompassBrowser,
   index: number

packages/compass-e2e-tests/tests/collection-ai-query.test.ts

@@ -45,7 +45,6 @@ describe('Collection ai query', function () {
     clearRequests = _clearRequests;
     setMockAtlasServerResponse = _setMockAtlasServerResponse;
 
-    process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE = endpoint;
     process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = endpoint;
 
     telemetry = await startTelemetryServer();
@@ -73,7 +72,6 @@ describe('Collection ai query', function () {
 
     await stopMockAtlasServer();
 
-    delete process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE;
     delete process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN;
 
     await cleanup(compass);
@@ -101,19 +99,16 @@ describe('Collection ai query', function () {
 
     it('makes request to the server and updates the query bar with the response', async function () {
       // Click the ai entry button.
-      await browser.clickVisible(Selectors.QueryBarAIEntryButton);
+      await browser.clickVisible(Selectors.GenAIEntryButton);
 
       // Enter the ai prompt.
-      await browser.clickVisible(Selectors.QueryBarAITextInput);
+      await browser.clickVisible(Selectors.GenAITextInput);
 
       const testUserInput = 'find all documents where i is greater than 50';
-      await browser.setValueVisible(
-        Selectors.QueryBarAITextInput,
-        testUserInput
-      );
+      await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
 
       // Click generate.
-      await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
+      await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
 
       // Wait for the ipc events to succeed.
       await browser.waitUntil(async function () {
@@ -163,22 +158,19 @@ describe('Collection ai query', function () {
 
     it('the error is shown to the user', async function () {
       // Click the ai entry button.
-      await browser.clickVisible(Selectors.QueryBarAIEntryButton);
+      await browser.clickVisible(Selectors.GenAIEntryButton);
 
       // Enter the ai prompt.
-      await browser.clickVisible(Selectors.QueryBarAITextInput);
+      await browser.clickVisible(Selectors.GenAITextInput);
 
       const testUserInput = 'find all documents where i is greater than 50';
-      await browser.setValueVisible(
-        Selectors.QueryBarAITextInput,
-        testUserInput
-      );
+      await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
 
       // Click generate.
-      await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
+      await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
 
       // Check that the error is shown.
-      const errorBanner = browser.$(Selectors.QueryBarAIErrorMessageBanner);
+      const errorBanner = browser.$(Selectors.GenAIErrorMessageBanner);
       await errorBanner.waitForDisplayed();
       expect(await errorBanner.getText()).to.equal(
         'Sorry, we were unable to generate the query, please try again. If the error persists, try changing your prompt.'

packages/compass-generative-ai/src/components/ai-experience-entry.tsx

@@ -86,7 +86,7 @@ const aiEntryLightModeStyles = css(
 );
 
 function AIExperienceEntry({
-  'data-testid': dataTestId,
+  'data-testid': dataTestId = 'open-gen-ai-button',
   type,
   onClick,
 }: {