chore(compass-e2e-tests): add generative query ai e2e tests with mock server COMPASS-6978

packages/atlas-service/src/components/ai-signin-modal.tsx

@@ -93,6 +93,7 @@ const AISignInModal: React.FunctionComponent<SignInModalProps> = ({
           <Button
             variant="primary"
             onClick={onSignInClick}
+            data-testid="atlas-signin-modal-button"
             className={buttonStyles}
             disabled={isSignInInProgress}
             // TODO: will have to update leafygreen for that

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

@@ -88,6 +88,7 @@ function createAIPlaceholderHTMLPlaceholder({
   containerEl.appendChild(placeholderTextEl);
 
   const aiButtonEl = document.createElement('button');
+  aiButtonEl.setAttribute('data-testid', 'open-ai-query-ask-ai-button');
   // By default placeholder container will have pointer events disabled
   aiButtonEl.style.pointerEvents = 'auto';
   // We stop mousedown from propagating and preventing default behavior to avoid

packages/compass-components/src/components/generative-ai/generative-ai-input.tsx

@@ -223,6 +223,7 @@ function GenerativeAIInput({
             className={textInputStyles}
             ref={promptTextInputRef}
             sizeVariant="small"
+            data-testid="ai-query-user-text-input"
             aria-label="Enter a plain text query that the AI will translate into MongoDB query language."
             placeholder="Tell Compass what documents to find (e.g. which movies were released in 2000)"
             value={aiPromptText}
@@ -247,6 +248,7 @@ function GenerativeAIInput({
                 generateButtonStyles,
                 !darkMode && generateButtonLightModeStyles
               )}
+              data-testid="ai-query-generate-button"
               onClick={() =>
                 isFetching ? onCancelRequest() : onSubmitText(aiPromptText)
               }
@@ -304,7 +306,9 @@ function GenerativeAIInput({
       </div>
       {errorMessage && (
         <div className={errorSummaryContainer}>
-          <ErrorSummary errors={errorMessage}>{errorMessage}</ErrorSummary>
+          <ErrorSummary data-testid="ai-query-error-msg" errors={errorMessage}>
+            {errorMessage}
+          </ErrorSummary>
         </div>
       )}
     </div>

packages/compass-e2e-tests/helpers/atlas-service.ts

@@ -0,0 +1,100 @@
+import http from 'http';
+import { once } from 'events';
+import type { AddressInfo } from 'net';
+
+export type MockAtlasServerResponse = {
+  status: number;
+  body: any;
+};
+
+export async function startMockAtlasServiceServer(
+  {
+    response: _response,
+  }: {
+    response: MockAtlasServerResponse;
+  } = {
+    response: {
+      status: 200,
+      body: {
+        content: {
+          query: {
+            filter: {
+              test: 'pineapple',
+            },
+          },
+        },
+      },
+    },
+  }
+): Promise<{
+  clearRequests: () => void;
+  getRequests: () => {
+    content: any;
+    req: any;
+  }[];
+  setMockAtlasServerResponse: (response: MockAtlasServerResponse) => void;
+  endpoint: string;
+  server: http.Server;
+  stop: () => Promise<void>;
+}> {
+  let requests: {
+    content: any;
+    req: any;
+  }[] = [];
+  let response = _response;
+  const server = http
+    .createServer((req, res) => {
+      let body = '';
+      req
+        .setEncoding('utf8')
+        .on('data', (chunk) => {
+          body += chunk;
+        })
+        .on('end', () => {
+          const jsonObject = JSON.parse(body);
+          requests.push({
+            req,
+            content: jsonObject,
+          });
+
+          res.setHeader('Content-Type', 'application/json');
+          if (response.status !== 200) {
+            res.writeHead(response.status);
+          }
+          return res.end(JSON.stringify(response.body));
+        });
+    })
+    .listen(0);
+  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 clearRequests() {
+    requests = [];
+  }
+
+  function getRequests() {
+    return requests;
+  }
+
+  function setMockAtlasServerResponse(newResponse: MockAtlasServerResponse) {
+    response = newResponse;
+  }
+
+  return {
+    clearRequests,
+    getRequests,
+    endpoint,
+    server,
+    setMockAtlasServerResponse,
+    stop,
+  };
+}

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

@@ -376,9 +376,6 @@ export const ShellExpandButton = '[data-testid="shell-expand-button"]';
 export const ShellInputEditor = '[data-testid="shell-input"] [data-codemirror]';
 export const ShellOutput = '[data-testid="shell-output"]';
 
-// Query bar (Find, Schema)
-export const QueryBarMenuActions = '#query-bar-menu-actions';
-
 // Instance screen
 export const InstanceTabs = '[data-testid="instance-tabs"]';
 export const InstanceTab = '.test-tab-nav-bar-tab';
@@ -1035,6 +1032,13 @@ export const queryBarExportToLanguageButton = (tabName: string): string => {
   const tabSelector = collectionContent(tabName);
   return `${tabSelector} [data-testid="query-bar-open-export-to-language-button"]`;
 };
+export const QueryBarAskAIButton =
+  '[data-testid="open-ai-query-ask-ai-button"]';
+export const QueryBarAITextInput = '[data-testid="ai-query-user-text-input"]';
+export const QueryBarAIGenerateQueryButton =
+  '[data-testid="ai-query-generate-button"]';
+export const QueryBarAIErrorMessageBanner =
+  '[data-testid="ai-query-error-msg"]';
 
 // Workspace tabs at the top
 export const SelectedWorkspaceTabButton =

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

@@ -0,0 +1,164 @@
+import chai from 'chai';
+
+import type { CompassBrowser } from '../helpers/compass-browser';
+import { startTelemetryServer } from '../helpers/telemetry';
+import type { Telemetry } from '../helpers/telemetry';
+import { beforeTests, afterTests, afterTest } from '../helpers/compass';
+import type { Compass } from '../helpers/compass';
+import * as Selectors from '../helpers/selectors';
+import { createNumbersCollection } from '../helpers/insert-data';
+import { startMockAtlasServiceServer } from '../helpers/atlas-service';
+import type { MockAtlasServerResponse } from '../helpers/atlas-service';
+import { getFirstListDocument } from '../helpers/read-first-document-content';
+
+const { expect } = chai;
+
+describe('Collection ai query', function () {
+  let compass: Compass;
+  let browser: CompassBrowser;
+  let telemetry: Telemetry;
+  let setMockAtlasServerResponse: (response: MockAtlasServerResponse) => void;
+  let stopMockAtlasServer: () => Promise<void>;
+  let getRequests: () => any[];
+  let clearRequests: () => void;
+
+  before(async function () {
+    process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN = 'true';
+
+    // Start a mock server to pass an ai response.
+    const {
+      endpoint,
+      getRequests: _getRequests,
+      clearRequests: _clearRequests,
+      setMockAtlasServerResponse: _setMockAtlasServerResponse,
+      stop,
+    } = await startMockAtlasServiceServer();
+
+    stopMockAtlasServer = stop;
+    getRequests = _getRequests;
+    clearRequests = _clearRequests;
+    setMockAtlasServerResponse = _setMockAtlasServerResponse;
+
+    process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE = endpoint;
+
+    telemetry = await startTelemetryServer();
+    compass = await beforeTests({
+      extraSpawnArgs: ['--enableAIExperience'],
+    });
+    browser = compass.browser;
+  });
+
+  beforeEach(async function () {
+    await createNumbersCollection();
+    await browser.connectWithConnectionString();
+    await browser.navigateToCollectionTab('test', 'numbers', 'Documents');
+  });
+
+  after(async function () {
+    await stopMockAtlasServer();
+
+    delete process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE;
+    delete process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN;
+
+    await afterTests(compass, this.currentTest);
+    await telemetry.stop();
+  });
+
+  afterEach(async function () {
+    clearRequests();
+    await afterTest(compass, this.currentTest);
+  });
+
+  describe('when the ai model response is valid', function () {
+    beforeEach(function () {
+      setMockAtlasServerResponse({
+        status: 200,
+        body: {
+          content: {
+            query: {
+              filter: {
+                i: {
+                  $gt: 50,
+                },
+              },
+            },
+          },
+        },
+      });
+    });
+
+    it('makes request to the server and updates the query bar with the response', async function () {
+      // Click the ask ai button.
+      await browser.clickVisible(Selectors.QueryBarAskAIButton);
+
+      // Enter the ai prompt.
+      await browser.clickVisible(Selectors.QueryBarAITextInput);
+
+      const testUserInput = 'find all documents where i is greater than 50';
+      await browser.setOrClearValue(
+        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 === '{i: {$gt: 50}}';
+      });
+
+      // Check that the request was made with the correct parameters.
+      const requests = getRequests();
+      expect(requests.length).to.equal(1);
+      expect(requests[0].content.userInput).to.equal(testUserInput);
+      expect(requests[0].content.collectionName).to.equal('numbers');
+      expect(requests[0].content.databaseName).to.equal('test');
+      expect(requests[0].content.schema).to.exist;
+
+      // Run it and check that the correct documents are shown.
+      await browser.runFind('Documents', true);
+      const modifiedResult = await getFirstListDocument(browser);
+      expect(modifiedResult.i).to.be.equal('51');
+    });
+  });
+
+  describe('when the Atlas service request errors', function () {
+    beforeEach(function () {
+      setMockAtlasServerResponse({
+        status: 500,
+        body: {
+          content: 'error',
+        },
+      });
+    });
+
+    it('the error is shown to the user', async function () {
+      // Click the ask ai button.
+      await browser.clickVisible(Selectors.QueryBarAskAIButton);
+
+      // Enter the ai prompt.
+      await browser.clickVisible(Selectors.QueryBarAITextInput);
+
+      const testUserInput = 'find all documents where i is greater than 50';
+      await browser.setOrClearValue(
+        Selectors.QueryBarAITextInput,
+        testUserInput
+      );
+
+      // Click generate.
+      await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
+
+      // Check that the error is shown.
+      const errorBanner = await browser.$(
+        Selectors.QueryBarAIErrorMessageBanner
+      );
+      await errorBanner.waitForDisplayed();
+      expect(await errorBanner.getText()).to.equal('500 Internal Server Error');
+    });
+  });
+});

packages/compass-query-bar/src/stores/ai-query-reducer.ts

@@ -270,7 +270,9 @@ export const cancelAIQuery = (): QueryBarThunkAction<
 export const showInput = (): QueryBarThunkAction<Promise<void>> => {
   return async (dispatch, _getState, { atlasService }) => {
     try {
-      await atlasService.signIn({ promptType: 'ai-promo-modal' });
+      if (process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN !== 'true') {
+        await atlasService.signIn({ promptType: 'ai-promo-modal' });
+      }
       dispatch({ type: AIQueryActionTypes.ShowInput });
     } catch {
       // if sign in failed / user canceled we just don't show the input