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

packages/compass-e2e-tests/helpers/atlas-ai-preferences-override.ts

@@ -0,0 +1,66 @@
+import http from 'http';
+import { once } from 'events';
+import type { AddressInfo } from 'net';
+
+export const E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT =
+  (process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT ??= '8081');
+
+export type PreferencesServerResponse = {
+  enableGenAIFeaturesAtlasProject: boolean;
+  enableGenAISampleDocumentPassingOnAtlasProject: boolean;
+  enableGenAIFeaturesAtlasOrg: boolean;
+  optInDataExplorerGenAIFeatures: boolean;
+};
+
+function preferencesResponse(
+  res: http.ServerResponse,
+  response: PreferencesServerResponse
+) {
+  // Get request to hello service.
+  res.setHeader('Content-Type', 'application/json');
+  return res.end(JSON.stringify(response));
+}
+
+export const enabledPreferencesResponse = {
+  enableGenAIFeaturesAtlasProject: true,
+  enableGenAISampleDocumentPassingOnAtlasProject: true,
+  enableGenAIFeaturesAtlasOrg: true,
+  optInDataExplorerGenAIFeatures: true,
+};
+
+export async function startPreferencesOverrideServer(
+  response: PreferencesServerResponse = enabledPreferencesResponse
+): Promise<{
+  setPreferencesResponse: (response: PreferencesServerResponse) => void;
+  endpoint: string;
+  server: http.Server;
+  stop: () => Promise<void>;
+}> {
+  const server = http
+    .createServer((req, res) => {
+      return preferencesResponse(res, response);
+    })
+    .listen(Number(E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT));
+  await once(server, 'listening');
+
+  // address() returns either a string or AddressInfo.
+  const address = server.address() as AddressInfo;
+
+  const endpoint = `http://localhost:${address.port}`;
+
+  async function stop() {
+    server.close();
+    await once(server, 'close');
+  }
+
+  function setPreferencesResponse(newResponse: PreferencesServerResponse) {
+    response = newResponse;
+  }
+
+  return {
+    endpoint,
+    server,
+    setPreferencesResponse,
+    stop,
+  };
+}

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

@@ -9,6 +9,7 @@ import { hideBin } from 'yargs/helpers';
 import Debug from 'debug';
 import fs from 'fs';
 import { getAtlasCloudSandboxDefaultConnections } from './compass-web-sandbox';
+import { E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT } from './atlas-ai-preferences-override';
 
 const debug = Debug('compass-e2e-tests:context');
 
@@ -344,6 +345,11 @@ process.env.HADRON_DISTRIBUTION ??= context.hadronDistribution;
 process.env.COMPASS_WEB_HTTP_PROXY_CLOUD_CONFIG ??=
   context.atlasCloudSandboxCloudConfig ?? 'dev';
 
+if (isTestingAtlasCloudSandbox(context)) {
+  process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT =
+    E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT;
+}
+
 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,140 @@
+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 {
+  enabledPreferencesResponse,
+  startPreferencesOverrideServer,
+} from '../../helpers/atlas-ai-preferences-override';
+import type { PreferencesServerResponse } from '../../helpers/atlas-ai-preferences-override';
+
+describe('Collection ai query', function () {
+  let compass: Compass;
+  let browser: CompassBrowser;
+  let setPreferencesResponse: (response: PreferencesServerResponse) => void;
+  let stopPreferencesServer: () => Promise<void>;
+
+  before(async function () {
+    if (!isTestingAtlasCloudSandbox()) {
+      this.skip();
+    }
+
+    // Start a mock server to pass an ai response.
+    const { setPreferencesResponse: _setPreferencesResponse, stop } =
+      await startPreferencesOverrideServer();
+
+    stopPreferencesServer = stop;
+    setPreferencesResponse = _setPreferencesResponse;
+  });
+
+  afterEach(async function () {
+    await screenshotIfFailed(compass, this.currentTest);
+    await cleanup(compass);
+  });
+
+  after(async function () {
+    if (!isTestingAtlasCloudSandbox()) {
+      return;
+    }
+
+    await stopPreferencesServer();
+  });
+
+  describe('when the feature is enabled', function () {
+    beforeEach(async function () {
+      setPreferencesResponse(enabledPreferencesResponse);
+
+      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'
+      );
+    });
+
+    it('should update the query bar with a generated query', async function () {
+      // Click the ai entry button.
+      await browser.clickVisible(Selectors.QueryBarAIEntryButton);
+
+      // Enter the ai prompt.
+      await browser.clickVisible(Selectors.QueryBarAITextInput);
+
+      const testUserInput = 'find all documents where i is greater than 50';
+      await browser.setValueVisible(
+        Selectors.QueryBarAITextInput,
+        testUserInput
+      );
+
+      // Click generate.
+      await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
+
+      // 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('when the org feature is disabled', function () {
+    beforeEach(async function () {
+      setPreferencesResponse({
+        enableGenAIFeaturesAtlasProject: true,
+        enableGenAISampleDocumentPassingOnAtlasProject: true,
+        enableGenAIFeaturesAtlasOrg: false,
+        optInDataExplorerGenAIFeatures: true,
+      });
+
+      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'
+      );
+    });
+
+    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.QueryBarAIEntryButton);
+      const isSidebarCreateCollectionButtonExisting =
+        await aiIntroButton.isExisting();
+      expect(isSidebarCreateCollectionButtonExisting).to.be.equal(false);
+    });
+  });
+});

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);

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

@@ -298,7 +298,39 @@ export class AtlasAiService {
     return body;
   }
 
+  async _e2eTestOnlyOverrideAIFeaturePreferences() {
+    try {
+      // When we're running e2e tests and want to customize these preferences on the fly
+      // we make a request to the server to override the preferences.
+      const {
+        enableGenAIFeaturesAtlasProject,
+        enableGenAISampleDocumentPassingOnAtlasProject,
+        enableGenAIFeaturesAtlasOrg,
+        optInDataExplorerGenAIFeatures,
+      } = await fetch('e2e-test-atlas-preferences-override').then((res) =>
+        res.json()
+      );
+
+      await this.preferences.savePreferences({
+        enableGenAIFeaturesAtlasProject,
+        enableGenAISampleDocumentPassingOnAtlasProject,
+        enableGenAIFeaturesAtlasOrg,
+        optInDataExplorerGenAIFeatures,
+      });
+      return;
+    } catch (e) {
+      /** no-op when the server isn't up. */
+    }
+  }
+
   async setupAIAccess(): Promise<void> {
+    if (
+      process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT !== undefined &&
+      process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT !== 'false'
+    ) {
+      await this._e2eTestOnlyOverrideAIFeaturePreferences();
+    }
+
     try {
       const featureResponse = await this.getAIFeatureEnablement();
 

packages/compass-preferences-model/src/compass-web-preferences-access.ts

@@ -15,6 +15,13 @@ export class CompassWebPreferencesAccess implements PreferencesAccess {
   }
 
   savePreferences(_attributes: Partial<UserPreferences>) {
+    if (
+      process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT !== undefined &&
+      process.env.E2E_TEST_ATLAS_PREFERENCES_OVERRIDE_PORT !== 'false'
+    ) {
+      return Promise.resolve(this._preferences.savePreferences(_attributes));
+    }
+
     // Only allow saving the optInDataExplorerGenAIFeatures preference.
     if (
       Object.keys(_attributes).length === 1 &&

packages/compass-web/sandbox/index.tsx

@@ -1,13 +1,19 @@
 import React, { useLayoutEffect } from 'react';
 import ReactDOM from 'react-dom';
-import { resetGlobalCSS, css, Body } from '@mongodb-js/compass-components';
+import {
+  resetGlobalCSS,
+  css,
+  Body,
+  SpinLoaderWithLabel,
+} from '@mongodb-js/compass-components';
 import { CompassWeb } from '../src/index';
 import { SandboxConnectionStorageProvider } from '../src/connection-storage';
 import { sandboxLogger } from './sandbox-logger';
 import { sandboxTelemetry } from './sandbox-telemetry';
 import { useAtlasProxySignIn } from './sandbox-atlas-sign-in';
 import { sandboxConnectionStorage } from './sandbox-connection-storage';
 import { useWorkspaceTabRouter } from './sandbox-workspace-tab-router';
+import { useAtlasPreferences } from './sandbox-atlas-preferences';
 
 const sandboxContainerStyles = css({
   width: '100%',
@@ -32,6 +38,10 @@ const App = () => {
   const [currentTab, updateCurrentTab] = useWorkspaceTabRouter();
   const { status, projectParams } = useAtlasProxySignIn();
   const { projectId, csrfToken, csrfTime } = projectParams ?? {};
+  const { status: preferencesStatus, preferences: atlasPreferences } =
+    useAtlasPreferences({
+      projectId,
+    });
 
   const atlasServiceSandboxBackendVariant =
     process.env.COMPASS_WEB_HTTP_PROXY_CLOUD_CONFIG === 'local'
@@ -41,9 +51,6 @@ const App = () => {
       ? 'web-sandbox-atlas-dev'
       : 'web-sandbox-atlas';
 
-  const overrideGenAIEnablement =
-    process.env.COMPASS_WEB_GEN_AI_ENABLEMENT === 'true';
-
   useLayoutEffect(() => {
     getMetaEl('csrf-token').setAttribute('content', csrfToken ?? '');
     getMetaEl('csrf-time').setAttribute('content', csrfTime ?? '');
@@ -55,6 +62,10 @@ const App = () => {
 
   const isAtlas = status === 'signed-in';
 
+  if (isAtlas && preferencesStatus === 'loading') {
+    return <SpinLoaderWithLabel progressText="Loading Atlas preferences" />;
+  }
+
   return (
     <SandboxConnectionStorageProvider
       value={isAtlas ? null : sandboxConnectionStorage}
@@ -80,11 +91,15 @@ const App = () => {
             enableCreatingNewConnections: !isAtlas,
             enableGlobalWrites: isAtlas,
             enableRollingIndexes: isAtlas,
-            enableGenAIFeaturesAtlasProject: isAtlas && overrideGenAIEnablement,
+            enableGenAIFeaturesAtlasProject:
+              isAtlas && !!atlasPreferences?.enableGenAIFeaturesAtlasProject,
             enableGenAISampleDocumentPassingOnAtlasProject:
-              isAtlas && overrideGenAIEnablement,
-            enableGenAIFeaturesAtlasOrg: isAtlas && overrideGenAIEnablement,
-            optInDataExplorerGenAIFeatures: isAtlas && overrideGenAIEnablement,
+              isAtlas &&
+              !!atlasPreferences?.enableGenAISampleDocumentPassingOnAtlasProject,
+            enableGenAIFeaturesAtlasOrg:
+              isAtlas && !!atlasPreferences?.enableGenAIFeaturesAtlasOrg,
+            optInDataExplorerGenAIFeatures:
+              isAtlas && !!atlasPreferences?.optInDataExplorerGenAIFeatures,
           }}
           onTrack={sandboxTelemetry.track}
           onDebug={sandboxLogger.debug}