chore: fixups

Author: gagik

packages/compass-e2e-tests/tests/assistant.test.ts

@@ -38,8 +38,6 @@ describe('MongoDB Assistant', function () {
   const collectionName = 'entryPoints';
 
   before(async function () {
-    skipForWeb(this, 'ai assistant not yet available in compass-web');
-
     process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN = 'true';
 
     // Start a mock Atlas service for feature flag checks
@@ -55,7 +53,6 @@ describe('MongoDB Assistant', function () {
 
     telemetry = await startTelemetryServer();
     compass = await init(this.test?.fullTitle());
-    browser = compass.browser;
 
     sendMessage = async (
       text: string,
@@ -75,34 +72,41 @@ describe('MongoDB Assistant', function () {
       await submitButton.click();
     };
 
+    const setup = async () => {
+      browser = compass.browser;
+      await browser.setupDefaultConnections();
+      await browser.connectToDefaults();
+      await browser.selectConnectionMenuItem(
+        DEFAULT_CONNECTION_NAME_1,
+        Selectors.CreateDatabaseButton,
+        false
+      );
+      await browser.addDatabase(dbName, collectionName);
+
+      await browser.navigateToCollectionTab(
+        DEFAULT_CONNECTION_NAME_1,
+        dbName,
+        collectionName,
+        'Aggregations'
+      );
+    };
+
     setAIOptIn = async (newValue: boolean) => {
       if (
         (await browser.getFeature('optInGenAIFeatures')) === true &&
         newValue === false
       ) {
+        await cleanup(compass);
         // Reseting the opt-in to false can be tricky so it's best to start over in this case.
-        compass = await init(this.test?.fullTitle(), { firstRun: false });
+        compass = await init(this.test?.fullTitle(), { firstRun: true });
+        await setup();
         return;
       }
 
       await browser.setFeature('optInGenAIFeatures', newValue);
     };
 
-    await browser.setupDefaultConnections();
-    await browser.connectToDefaults();
-    await browser.selectConnectionMenuItem(
-      DEFAULT_CONNECTION_NAME_1,
-      Selectors.CreateDatabaseButton,
-      false
-    );
-    await browser.addDatabase(dbName, collectionName);
-
-    await browser.navigateToCollectionTab(
-      DEFAULT_CONNECTION_NAME_1,
-      dbName,
-      collectionName,
-      'Aggregations'
-    );
+    await setup();
   });
 
   after(async function () {
@@ -164,72 +168,6 @@ describe('MongoDB Assistant', function () {
     });
   });
 
-  describe('before opt-in', function () {
-    before(async function () {
-      await setAIOptIn(false);
-    });
-
-    it('does not send the message if the user declines the opt-in', async function () {
-      await openAssistantDrawer(browser);
-
-      await sendMessage(testMessage);
-
-      const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
-      await declineLink.waitForDisplayed();
-      await declineLink.click();
-
-      const optInModal = browser.$(Selectors.AIOptInModal);
-      await optInModal.waitForDisplayed({ reverse: true });
-
-      const chatInput = browser.$(Selectors.AssistantChatInputTextArea);
-      expect(await chatInput.getValue()).not.to.equal(testMessage);
-
-      expect(await getDisplayedMessages(browser)).to.deep.equal([]);
-
-      expect(mockAssistantServer.getRequests()).to.be.empty;
-    });
-
-    describe('entry points', function () {
-      it('should display opt-in modal for connection error entry point', async function () {
-        await browser.connectWithConnectionString(
-          'mongodb-invalid://localhost:27017',
-          { connectionStatus: 'failure' }
-        );
-        await browser.clickVisible(
-          browser.$(Selectors.ConnectionToastErrorDebugButton)
-        );
-
-        const optInModal = browser.$(Selectors.AIOptInModal);
-        await optInModal.waitForDisplayed();
-        expect(await optInModal.isDisplayed()).to.be.true;
-
-        const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
-        await declineLink.waitForDisplayed();
-        await declineLink.click();
-
-        await optInModal.waitForDisplayed({ reverse: true });
-
-        expect(await getDisplayedMessages(browser)).to.deep.equal([]);
-      });
-
-      it('should display opt-in modal for explain plan entry point', async function () {
-        await useExplainPlanEntryPoint(browser);
-
-        const optInModal = browser.$(Selectors.AIOptInModal);
-        await optInModal.waitForDisplayed();
-        expect(await optInModal.isDisplayed()).to.be.true;
-
-        const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
-        await declineLink.waitForDisplayed();
-        await declineLink.click();
-
-        await optInModal.waitForDisplayed({ reverse: true });
-
-        expect(await getDisplayedMessages(browser)).to.deep.equal([]);
-      });
-    });
-  });
-
   describe('opting in', function () {
     before(async function () {
       await setAIOptIn(false);
@@ -465,6 +403,72 @@ describe('MongoDB Assistant', function () {
       });
     });
   });
+
+  describe('before opt-in', function () {
+    before(async function () {
+      await setAIOptIn(false);
+    });
+
+    it('does not send the message if the user declines the opt-in', async function () {
+      await openAssistantDrawer(browser);
+
+      await sendMessage(testMessage);
+
+      const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
+      await declineLink.waitForDisplayed();
+      await declineLink.click();
+
+      const optInModal = browser.$(Selectors.AIOptInModal);
+      await optInModal.waitForDisplayed({ reverse: true });
+
+      const chatInput = browser.$(Selectors.AssistantChatInputTextArea);
+      expect(await chatInput.getValue()).not.to.equal(testMessage);
+
+      expect(await getDisplayedMessages(browser)).to.deep.equal([]);
+
+      expect(mockAssistantServer.getRequests()).to.be.empty;
+    });
+
+    describe('entry points', function () {
+      it('should display opt-in modal for connection error entry point', async function () {
+        await browser.connectWithConnectionString(
+          'mongodb-invalid://localhost:27017',
+          { connectionStatus: 'failure' }
+        );
+        await browser.clickVisible(
+          browser.$(Selectors.ConnectionToastErrorDebugButton)
+        );
+
+        const optInModal = browser.$(Selectors.AIOptInModal);
+        await optInModal.waitForDisplayed();
+        expect(await optInModal.isDisplayed()).to.be.true;
+
+        const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
+        await declineLink.waitForDisplayed();
+        await declineLink.click();
+
+        await optInModal.waitForDisplayed({ reverse: true });
+
+        expect(await getDisplayedMessages(browser)).to.deep.equal([]);
+      });
+
+      it('should display opt-in modal for explain plan entry point', async function () {
+        await useExplainPlanEntryPoint(browser);
+
+        const optInModal = browser.$(Selectors.AIOptInModal);
+        await optInModal.waitForDisplayed();
+        expect(await optInModal.isDisplayed()).to.be.true;
+
+        const declineLink = browser.$(Selectors.AIOptInModalDeclineLink);
+        await declineLink.waitForDisplayed();
+        await declineLink.click();
+
+        await optInModal.waitForDisplayed({ reverse: true });
+
+        expect(await getDisplayedMessages(browser)).to.deep.equal([]);
+      });
+    });
+  });
 });
 
 async function setAIFeatures(browser: CompassBrowser, newValue: boolean) {