Skip to content

Commit 44e5585

Browse files
gagikgagik
committed
chore: use a feature flag
1 parent 3df8f4a commit 44e5585

File tree

4 files changed

+363
-1
lines changed

4 files changed

+363
-1
lines changed

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

Lines changed: 336 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,336 @@
1+
import type { CompassBrowser } from '../helpers/compass-browser';
2+
import {
3+
init,
4+
cleanup,
5+
screenshotIfFailed,
6+
Selectors,
7+
skipForWeb,
8+
TEST_COMPASS_WEB,
9+
DEFAULT_CONNECTION_NAME_1,
10+
} from '../helpers/compass';
11+
import type { Compass } from '../helpers/compass';
12+
import type { OIDCMockProviderConfig } from '@mongodb-js/oidc-mock-provider';
13+
import { OIDCMockProvider } from '@mongodb-js/oidc-mock-provider';
14+
import path from 'path';
15+
import { expect } from 'chai';
16+
import { createNumbersCollection } from '../helpers/insert-data';
17+
import { startMockAtlasServiceServer } from '../helpers/atlas-service';
18+
import type { Telemetry } from '../helpers/telemetry';
19+
import { startTelemetryServer } from '../helpers/telemetry';
20+
21+
const DEFAULT_TOKEN_PAYLOAD = {
22+
expires_in: 3600,
23+
payload: {
24+
groups: ['testgroup'],
25+
sub: 'testuser',
26+
aud: 'resource-server-audience-value',
27+
},
28+
};
29+
30+
function getTestBrowserShellCommand() {
31+
return `${process.execPath} ${path.resolve(
32+
__dirname,
33+
'..',
34+
'fixtures',
35+
'curl.js'
36+
)}`;
37+
}
38+
39+
describe('Atlas Login', function () {
40+
let compass: Compass;
41+
let browser: CompassBrowser;
42+
let oidcMockProvider: OIDCMockProvider;
43+
let getTokenPayload: OIDCMockProviderConfig['getTokenPayload'];
44+
let stopMockAtlasServer: () => Promise<void>;
45+
let numberOfOIDCAuthRequests = 0;
46+
47+
before(async function () {
48+
skipForWeb(this, 'atlas-login not supported in compass-web');
49+
50+
// Start a mock server to pass an ai response.
51+
const { endpoint, stop } = await startMockAtlasServiceServer();
52+
stopMockAtlasServer = stop;
53+
process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = endpoint;
54+
55+
function isAuthorised(req: { headers: { authorization?: string } }) {
56+
const [, token] = req.headers.authorization?.split(' ') ?? [];
57+
// We can't check that the issued token is the one received by oidc-plugin
58+
// so we are only checking that it was passed and assuming that it's good
59+
// enough of a validation
60+
return !!token;
61+
}
62+
63+
oidcMockProvider = await OIDCMockProvider.create({
64+
getTokenPayload(metadata) {
65+
return getTokenPayload(metadata);
66+
},
67+
overrideRequestHandler(_url, req, res) {
68+
const url = new URL(_url);
69+
70+
switch (url.pathname) {
71+
case '/auth-portal-redirect':
72+
res.statusCode = 307;
73+
res.setHeader('Location', url.searchParams.get('fromURI') ?? '');
74+
res.end();
75+
break;
76+
case '/authorize':
77+
numberOfOIDCAuthRequests += 1;
78+
break;
79+
case '/v1/userinfo':
80+
if (isAuthorised(req)) {
81+
res.statusCode = 200;
82+
res.write(
83+
JSON.stringify({
84+
sub: Date.now().toString(32),
85+
firstName: 'First',
86+
lastName: 'Last',
87+
primaryEmail: '[email protected]',
88+
login: '[email protected]',
89+
})
90+
);
91+
res.end();
92+
} else {
93+
res.statusCode = 401;
94+
res.end();
95+
}
96+
break;
97+
case '/v1/introspect':
98+
res.statusCode = 200;
99+
res.write(JSON.stringify({ active: isAuthorised(req) }));
100+
res.end();
101+
break;
102+
}
103+
},
104+
});
105+
106+
process.env.COMPASS_CLIENT_ID_OVERRIDE = 'testServer';
107+
process.env.COMPASS_OIDC_ISSUER_OVERRIDE = oidcMockProvider.issuer;
108+
process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE = `${oidcMockProvider.issuer}/auth-portal-redirect`;
109+
});
110+
111+
beforeEach(async function () {
112+
numberOfOIDCAuthRequests = 0;
113+
114+
getTokenPayload = () => {
115+
return DEFAULT_TOKEN_PAYLOAD;
116+
};
117+
118+
compass = await init(this.test?.fullTitle());
119+
browser = compass.browser;
120+
await browser.setFeature(
121+
'browserCommandForOIDCAuth',
122+
getTestBrowserShellCommand()
123+
);
124+
await browser.setupDefaultConnections();
125+
});
126+
127+
afterEach(async function () {
128+
await browser.setFeature('browserCommandForOIDCAuth', undefined);
129+
await screenshotIfFailed(compass, this.currentTest);
130+
await cleanup(compass);
131+
});
132+
133+
after(async function () {
134+
if (TEST_COMPASS_WEB) {
135+
return;
136+
}
137+
138+
await oidcMockProvider?.close();
139+
delete process.env.COMPASS_CLIENT_ID_OVERRIDE;
140+
delete process.env.COMPASS_OIDC_ISSUER_OVERRIDE;
141+
delete process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE;
142+
143+
await stopMockAtlasServer();
144+
delete process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE;
145+
});
146+
147+
describe('in settings', function () {
148+
it('should sign in user when clicking on "Log in with Atlas" button', async function () {
149+
await browser.openSettingsModal('ai');
150+
151+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
152+
153+
const loginStatus = browser.$(Selectors.AtlasLoginStatus);
154+
await browser.waitUntil(async () => {
155+
return (
156+
(await loginStatus.getText()).trim() ===
157+
'Logged in with Atlas account [email protected]'
158+
);
159+
});
160+
expect(numberOfOIDCAuthRequests).to.eq(1);
161+
});
162+
163+
describe('telemetry', () => {
164+
let telemetry: Telemetry;
165+
166+
before(async function () {
167+
telemetry = await startTelemetryServer();
168+
});
169+
170+
after(async function () {
171+
await telemetry.stop();
172+
});
173+
174+
it('should send identify after the user has logged in', async function () {
175+
const atlasUserIdBefore = await browser.getFeature(
176+
'telemetryAtlasUserId'
177+
);
178+
expect(atlasUserIdBefore).to.not.exist;
179+
180+
await browser.openSettingsModal('ai');
181+
182+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
183+
184+
const loginStatus = browser.$(Selectors.AtlasLoginStatus);
185+
await browser.waitUntil(async () => {
186+
return (
187+
(await loginStatus.getText()).trim() ===
188+
'Logged in with Atlas account [email protected]'
189+
);
190+
});
191+
192+
const atlasUserIdAfter = await browser.getFeature(
193+
'telemetryAtlasUserId'
194+
);
195+
expect(atlasUserIdAfter).to.be.a('string');
196+
197+
const identify = telemetry
198+
.events()
199+
.find((entry) => entry.type === 'identify');
200+
expect(identify.traits.platform).to.equal(process.platform);
201+
expect(identify.traits.arch).to.match(/^(x64|arm64)$/);
202+
});
203+
});
204+
205+
it('should sign out user when "Disconnect" clicked', async function () {
206+
await browser.openSettingsModal('ai');
207+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
208+
209+
const loginStatus = browser.$(Selectors.AtlasLoginStatus);
210+
211+
await browser.waitUntil(async () => {
212+
return (
213+
(await loginStatus.getText()).trim() ===
214+
'Logged in with Atlas account [email protected]'
215+
);
216+
});
217+
218+
await browser.clickVisible(Selectors.DisconnectAtlasAccountButton);
219+
220+
await browser.waitUntil(async () => {
221+
return (await loginStatus.getText()).includes(
222+
'This is a feature powered by generative AI, and may give inaccurate responses'
223+
);
224+
});
225+
});
226+
227+
it('should sign in user when disconnected and clicking again on "Log in with Atlas" button', async function () {
228+
await browser.openSettingsModal('ai');
229+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
230+
231+
let loginStatus = browser.$(Selectors.AtlasLoginStatus);
232+
233+
await browser.waitUntil(async () => {
234+
return (
235+
(await loginStatus.getText()).trim() ===
236+
'Logged in with Atlas account [email protected]'
237+
);
238+
});
239+
240+
await browser.clickVisible(Selectors.DisconnectAtlasAccountButton);
241+
242+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
243+
244+
loginStatus = browser.$(Selectors.AtlasLoginStatus);
245+
await browser.waitUntil(async () => {
246+
return (
247+
(await loginStatus.getText()).trim() ===
248+
'Logged in with Atlas account [email protected]'
249+
);
250+
});
251+
expect(numberOfOIDCAuthRequests).to.eq(2);
252+
});
253+
254+
it('should show toast with error if sign in failed', async function () {
255+
getTokenPayload = () => {
256+
return Promise.reject(new Error('Auth failed'));
257+
};
258+
259+
await browser.openSettingsModal('ai');
260+
await browser.clickVisible(Selectors.LogInWithAtlasButton);
261+
262+
const errorToast = browser.$(Selectors.AtlasLoginErrorToast);
263+
await errorToast.waitForDisplayed();
264+
265+
expect(await errorToast.getText()).to.match(
266+
/Sign in failed\n+unexpected HTTP response status code.+Auth failed/
267+
);
268+
});
269+
});
270+
271+
describe('in CRUD view', function () {
272+
beforeEach(async function () {
273+
await createNumbersCollection();
274+
await browser.disconnectAll();
275+
await browser.connectToDefaults();
276+
await browser.navigateToCollectionTab(
277+
DEFAULT_CONNECTION_NAME_1,
278+
'test',
279+
'numbers',
280+
'Documents'
281+
);
282+
});
283+
284+
it('should not show AI input if sign in flow was not finished', async function () {
285+
getTokenPayload = () => {
286+
return new Promise(() => {});
287+
};
288+
289+
const generateQueryButton = browser.$('button*=Generate query');
290+
await browser.clickVisible(generateQueryButton);
291+
292+
await browser.clickVisible(Selectors.LogInWithAtlasModalButton);
293+
294+
// Because leafygreen doesn't render a button there and we don't have any
295+
// control over it
296+
await browser.clickVisible('span=Not now');
297+
298+
const aiInput = browser.$(Selectors.GenAITextInput);
299+
expect(await aiInput.isExisting()).to.eq(false);
300+
expect(await generateQueryButton.isDisplayed()).to.eq(true);
301+
});
302+
});
303+
304+
describe('in Aggregation Builder view', function () {
305+
beforeEach(async function () {
306+
await createNumbersCollection();
307+
await browser.disconnectAll();
308+
await browser.connectToDefaults();
309+
await browser.navigateToCollectionTab(
310+
DEFAULT_CONNECTION_NAME_1,
311+
'test',
312+
'numbers',
313+
'Aggregations'
314+
);
315+
});
316+
317+
it('should not show AI input if sign in flow was not finished', async function () {
318+
getTokenPayload = () => {
319+
return new Promise(() => {});
320+
};
321+
322+
const generateQueryButton = browser.$('button*=Generate aggregation');
323+
await browser.clickVisible(generateQueryButton);
324+
325+
await browser.clickVisible(Selectors.LogInWithAtlasModalButton);
326+
327+
// Because leafygreen doesn't render a button there and we don't have any
328+
// control over it
329+
await browser.clickVisible('span=Not now');
330+
331+
const aiInput = browser.$(Selectors.GenAITextInput);
332+
expect(await aiInput.isExisting()).to.eq(false);
333+
expect(await generateQueryButton.isDisplayed()).to.eq(true);
334+
});
335+
});
336+
});

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,13 @@ describe('AtlasAiService', function () {
327327
// Reset preferences
328328
await preferences.savePreferences({
329329
optInGenAIFeatures: false,
330+
enableUnauthenticatedGenAI: true,
331+
});
332+
});
333+
334+
afterEach(async function () {
335+
await preferences.savePreferences({
336+
enableUnauthenticatedGenAI: false,
330337
});
331338
});
332339

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import type { Logger } from '@mongodb-js/compass-logging';
1111
import { EJSON } from 'bson';
1212
import { getStore } from './store/atlas-ai-store';
1313
import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer';
14+
import { signIntoAtlasWithModalPrompt } from './store/atlas-signin-reducer';
1415

1516
type GenerativeAiInput = {
1617
userInput: string;
@@ -276,7 +277,17 @@ export class AtlasAiService {
276277
}
277278

278279
async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) {
279-
return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal }));
280+
if (this.preferences.getPreferences().enableUnauthenticatedGenAI) {
281+
return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal }));
282+
}
283+
284+
// When the ai feature is attempted to be opened we make sure
285+
// the user is signed into Atlas and opted in.
286+
287+
if (this.apiURLPreset === 'cloud') {
288+
return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal }));
289+
}
290+
return getStore().dispatch(signIntoAtlasWithModalPrompt({ signal }));
280291
}
281292

282293
private getQueryOrAggregationFromUserInput = async <T>(

‎packages/compass-preferences-model/src/feature-flags.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ export type FeatureFlags = {
2828
showIndexesGuidanceVariant: boolean;
2929
enableContextMenus: boolean;
3030
enableSearchActivationProgramP1: boolean;
31+
enableUnauthenticatedGenAI: boolean;
3132
};
3233

3334
export const featureFlags: Required<{
@@ -150,6 +151,13 @@ export const featureFlags: Required<{
150151
},
151152
},
152153

154+
enableUnauthenticatedGenAI: {
155+
stage: 'development',
156+
description: {
157+
short: 'Enable GenAI for unauthenticated users',
158+
},
159+
},
160+
153161
/**
154162
* Feature flag for CLOUDP-308952.
155163
*/

0 commit comments

Comments
 (0)