chore(e2e): clean-up and re-enable genai tests for web (#7555)

* chore(atlas-service): resolve config in runtime to make sure that env overrides are applied

* chore(generative-ai): use atlasMetadata as a way to detect which url preset to use

* chore(assistant): dynamically resolve service url

* chore(e2e, web): remove override inlining from the build; make atlas service mock handle CORS properly; re-enable disabled genai tests

* chore(e2e): restore lost screenshots

Author: gribnoysup

packages/atlas-service/src/atlas-service.spec.ts

@@ -29,9 +29,10 @@ function getAtlasService(
   const atlasService = new AtlasService(
     authService,
     preferences,
-    createNoopLogger()
+    createNoopLogger(),
+    undefined,
+    ATLAS_CONFIG
   );
-  atlasService['config'] = ATLAS_CONFIG;
   return atlasService;
 }
 

packages/atlas-service/src/atlas-service.ts

@@ -53,14 +53,17 @@ function getAutomationAgentClusterId(
 }
 
 export class AtlasService {
-  private config: AtlasServiceConfig;
   constructor(
     private readonly authService: AtlasAuthService,
     private readonly preferences: PreferencesAccess,
     private readonly logger: Logger,
-    private readonly options?: AtlasServiceOptions
-  ) {
-    this.config = getAtlasConfig(preferences);
+    private readonly options?: AtlasServiceOptions,
+    private readonly defaultConfigOverride?: AtlasServiceConfig
+  ) {}
+  // Config value is dynamic to make sure that process.env overrides are taken
+  // into account in runtime
+  get config(): AtlasServiceConfig {
+    return this.defaultConfigOverride ?? getAtlasConfig(this.preferences);
   }
   adminApiEndpoint(path?: string): string {
     return `${this.config.atlasApiBaseUrl}${normalizePath(path)}`;

packages/atlas-service/src/util.ts

@@ -255,6 +255,7 @@ export function getAtlasConfig(
   const { atlasServiceBackendPreset } = preferences.getPreferences();
   const envConfig = {
     atlasApiBaseUrl: process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE,
+    cloudBaseUrl: process.env.COMPASS_CLOUD_BASE_URL_OVERRIDE,
     atlasLogin: {
       clientId: process.env.COMPASS_CLIENT_ID_OVERRIDE,
       issuer: process.env.COMPASS_OIDC_ISSUER_OVERRIDE,

packages/compass-assistant/src/compass-assistant-provider.spec.tsx

@@ -28,7 +28,7 @@ import { CompassAssistantDrawer } from './compass-assistant-drawer';
 import { createBrokenTransport, createMockChat } from '../test/utils';
 import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
 import type { TrackFunction } from '@mongodb-js/compass-telemetry';
-import { createLogger } from '@mongodb-js/compass-logging';
+import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
 
 function createMockProvider({
   mockAtlasService,
@@ -480,7 +480,7 @@ describe('CompassAssistantProvider', function () {
               .stub()
               .returns('https://localhost:3000'),
           } as unknown as AtlasService,
-          logger: createLogger('COMPASS-ASSISTANT-TEST'),
+          logger: createNoopLogger(),
           track: track as unknown as TrackFunction,
         });
         await renderOpenAssistantDrawer({
@@ -654,47 +654,42 @@ describe('CompassAssistantProvider', function () {
     });
   });
 
+  let sandbox: sinon.SinonSandbox;
+
+  beforeEach(function () {
+    sandbox = sinon.createSandbox();
+  });
+
+  afterEach(function () {
+    if (sandbox) {
+      sandbox.reset();
+    }
+  });
+
   it('uses the Atlas Service assistantApiEndpoint', async function () {
     const mockAtlasService = {
-      assistantApiEndpoint: sinon
+      assistantApiEndpoint: sandbox
         .stub()
         .returns('https://example.com/assistant/api/v1'),
     };
 
-    const mockAtlasAiService = {
-      ensureAiFeatureAccess: sinon.stub().callsFake(() => {
-        return Promise.resolve();
-      }),
-    };
-
-    const mockAtlasAuthService = {};
-
-    const MockedProvider = CompassAssistantProvider.withMockServices({
+    const chat = createDefaultChat({
+      originForPrompt: 'foo',
+      appNameForPrompt: 'bar',
       atlasService: mockAtlasService as unknown as AtlasService,
-      atlasAiService: mockAtlasAiService as unknown as AtlasAiService,
-      atlasAuthService: mockAtlasAuthService as unknown as AtlasAuthService,
+      logger: createNoopLogger(),
+      track: () => undefined,
     });
 
-    render(
-      <DrawerContentProvider>
-        <DrawerAnchor />
-        <MockedProvider
-          originForPrompt="mongodb-compass"
-          appNameForPrompt="MongoDB Compass"
-        />
-      </DrawerContentProvider>,
-      {
-        preferences: {
-          enableAIAssistant: true,
-          enableGenAIFeatures: true,
-          enableGenAIFeaturesAtlasOrg: true,
-          cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
-        },
-      }
-    );
+    const fetchStub = sandbox
+      .stub(globalThis, 'fetch')
+      .resolves({ ok: true, headers: [] } as any);
 
-    await waitFor(() => {
-      expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true;
-    });
+    await chat.sendMessage({ text: 'hello' });
+
+    expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true;
+    expect(fetchStub.lastCall.args[0]).to.eq(
+      'https://example.com/assistant/api/v1/responses'
+    );
   });
 });

packages/compass-assistant/src/compass-assistant-provider.tsx

@@ -334,6 +334,7 @@ export function createDefaultChat({
     transport: Chat<AssistantMessage>['transport'];
   };
 }): Chat<AssistantMessage> {
+  const initialBaseUrl = 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid';
   return new Chat({
     transport:
       options?.transport ??
@@ -343,8 +344,21 @@ export function createDefaultChat({
           target: appNameForPrompt,
         }),
         model: createOpenAI({
-          baseURL: atlasService.assistantApiEndpoint(),
+          baseURL: initialBaseUrl,
           apiKey: '',
+          fetch(url, init) {
+            return globalThis.fetch(
+              // The `baseUrl` can be dynamically changed, but `createOpenAI`
+              // constructor doesn't allow us to change it after initial call.
+              // Instead we're going to update it every time the fetch call
+              // happens
+              String(url).replace(
+                initialBaseUrl,
+                atlasService.assistantApiEndpoint()
+              ),
+              init
+            );
+          },
         }).responses('mongodb-chat-latest'),
       }),
     onError: (err: Error) => {

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

@@ -140,7 +140,6 @@ function sendStreamingResponse(res: http.ServerResponse, content: string) {
   sendChunk();
 }
 
-export const MOCK_ASSISTANT_SERVER_PORT = 27097;
 export async function startMockAssistantServer(
   {
     response: _response,
@@ -222,7 +221,7 @@ export async function startMockAssistantServer(
           return sendStreamingResponse(res, response.body);
         });
     })
-    .listen(MOCK_ASSISTANT_SERVER_PORT);
+    .listen(0);
   await once(server, 'listening');
 
   // address() returns either a string or AddressInfo.

packages/compass-e2e-tests/helpers/compass-web-sandbox.ts

@@ -11,7 +11,6 @@ import {
 } from './test-runner-paths';
 import type { ConnectionInfo } from '@mongodb-js/connection-info';
 import ConnectionString from 'mongodb-connection-string-url';
-import { MOCK_ASSISTANT_SERVER_PORT } from './assistant-service';
 
 const debug = Debug('compass-e2e-tests:compass-web-sandbox');
 
@@ -22,9 +21,6 @@ const debug = Debug('compass-e2e-tests:compass-web-sandbox');
 process.env.OPEN_BROWSER = 'false'; // tell webpack dev server not to open the default browser
 process.env.DISABLE_DEVSERVER_OVERLAY = 'true';
 process.env.APP_ENV = 'webdriverio';
-// Set the assistant base URL override for tests so we can mock the assistant server
-process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE = `http://localhost:${MOCK_ASSISTANT_SERVER_PORT}`;
-process.env.COMPASS_OVERRIDE_ENABLE_AI_FEATURES = 'true';
 
 const wait = (ms: number) => {
   return new Promise((resolve) => {

packages/compass-e2e-tests/helpers/atlas-service.ts renamed to packages/compass-e2e-tests/helpers/mock-atlas-service.ts

@@ -61,33 +61,62 @@ export async function startMockAtlasServiceServer(
   let response = _response;
   const server = http
     .createServer((req, res) => {
-      if (req.method === 'GET') {
-        requests.push({
-          req,
-          content: null,
-        });
-        return aiFeatureEnableResponse(req, res);
+      res.setHeader('Access-Control-Allow-Origin', req.headers.origin ?? '*');
+      res.setHeader('Access-Control-Allow-Credentials', 'true');
+      res.setHeader(
+        'Access-Control-Allow-Headers',
+        req.headers['access-control-request-headers'] ?? '*'
+      );
+
+      if (req.method === 'OPTIONS') {
+        res.statusCode = 200;
+        res.end();
+        return;
+      }
+
+      if (req.url?.startsWith('/settings/optInDataExplorerGenAIFeatures')) {
+        res.statusCode = 200;
+        res.end();
+        return;
       }
 
-      let body = '';
-      req
-        .setEncoding('utf8')
-        .on('data', (chunk) => {
-          body += chunk;
-        })
-        .on('end', () => {
-          const jsonObject = JSON.parse(body);
+      if (req.url?.startsWith('/unauth/ai/api/v1/mql-query')) {
+        if (req.method === 'GET') {
           requests.push({
             req,
-            content: jsonObject,
+            content: null,
           });
+          return aiFeatureEnableResponse(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));
+          });
+
+        return;
+      }
 
-          res.setHeader('Content-Type', 'application/json');
-          if (response.status !== 200) {
-            res.writeHead(response.status);
-          }
-          return res.end(JSON.stringify(response.body));
-        });
+      res.statusCode = 404;
+      res.statusMessage =
+        'Route not found in the mock Atlas backend. Did you forget to add a mock route support?';
+      res.end();
     })
     .listen(0);
   await once(server, 'listening');