Skip to content

Commit b13cd9b

Browse files
AnemyAnemy
authored
chore(e2e-tests): add gen ai test with logged in atlas user COMPASS-8395 (#6547)
1 parent 310b261 commit b13cd9b

File tree

17 files changed

+450
-116
lines changed

17 files changed

+450
-116
lines changed

‎packages/compass-e2e-tests/helpers/commands/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ export * from './unhide-index';
6262
export * from './hide-visible-modal';
6363
export * from './hide-visible-toasts';
6464
export * from './sidebar-collection';
65+
export * from './switch-pipeline-mode';
6566
export * from './read-first-document-content';
6667
export * from './read-stage-operators';
6768
export * from './click-confirmation-action';

‎packages/compass-e2e-tests/helpers/commands/set-feature.ts‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,34 @@
11
import type { CompassBrowser } from '../compass-browser';
2-
import type { UserPreferences } from 'compass-preferences-model';
2+
import type {
3+
AllPreferences,
4+
UserPreferences,
5+
} from 'compass-preferences-model';
6+
import { isTestingWeb } from '../test-runner-context';
37

48
export async function setFeature<K extends keyof UserPreferences>(
59
browser: CompassBrowser,
610
name: K,
711
value: UserPreferences[K]
812
): Promise<void> {
13+
if (isTestingWeb()) {
14+
// When running in Compass web we cannot use the IPC to update the
15+
// preferences so we use a global function.
16+
await browser.execute(
17+
async (_name, _value) => {
18+
const attributes: Partial<AllPreferences> = {
19+
[_name]: _value === null ? undefined : _value,
20+
};
21+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
22+
await (globalThis as any).__compassWebE2ETestSavePreferences(
23+
attributes
24+
);
25+
},
26+
name,
27+
value
28+
);
29+
return;
30+
}
31+
932
await browser.execute(
1033
async (_name, _value) => {
1134
// eslint-disable-next-line @typescript-eslint/no-var-requires

‎packages/compass-e2e-tests/helpers/commands/switch-pipeline-mode.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import type { CompassBrowser } from '../compass-browser';
2+
import * as Selectors from '../selectors';
3+
4+
export async function switchPipelineMode(
5+
browser: CompassBrowser,
6+
mode: 'as-text' | 'builder-ui'
7+
): Promise<void> {
8+
await browser.clickVisible(Selectors.aggregationPipelineModeToggle(mode));
9+
await browser.waitForAnimations(Selectors.AggregationBuilderWorkspace);
10+
}

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,12 +1239,11 @@ export const queryBarExportToLanguageButton = (tabName: string): string => {
12391239
const tabSelector = collectionContent(tabName);
12401240
return `${tabSelector} [data-testid="query-bar-open-export-to-language-button"]`;
12411241
};
1242-
export const QueryBarAIEntryButton =
1243-
'[data-testid="open-ai-query-entry-button"]';
1244-
export const QueryBarAITextInput = '[data-testid="ai-user-text-input"]';
1245-
export const QueryBarAIGenerateQueryButton =
1246-
'[data-testid="ai-generate-button"]';
1247-
export const QueryBarAIErrorMessageBanner = '[data-testid="ai-error-msg"]';
1242+
export const GenAIEntryButton = '[data-testid="open-ai-query-entry-button"]';
1243+
export const GenAITextInput = '[data-testid="ai-user-text-input"]';
1244+
export const GenAIGenerateQueryButton = '[data-testid="ai-generate-button"]';
1245+
export const GenAIErrorMessageBanner = '[data-testid="ai-error-msg"]';
1246+
export const GenAIOpenButton = '[data-testid="open-gen-ai-button"]';
12481247

12491248
// Workspace tabs
12501249
export const WorkspaceTabsContainer =

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,10 @@ process.env.HADRON_DISTRIBUTION ??= context.hadronDistribution;
344344
process.env.COMPASS_WEB_HTTP_PROXY_CLOUD_CONFIG ??=
345345
context.atlasCloudSandboxCloudConfig ?? 'dev';
346346

347+
if (isTestingAtlasCloudSandbox(context)) {
348+
process.env.E2E_TEST_CLOUD_WEB_ENABLE_PREFERENCE_SAVING ??= 'true';
349+
}
350+
347351
const testServerVersion =
348352
process.env.MONGODB_VERSION ?? process.env.MONGODB_RUNNER_VERSION;
349353

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

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
import { expect } from 'chai';
2+
3+
import type { CompassBrowser } from '../../helpers/compass-browser';
4+
import {
5+
init,
6+
cleanup,
7+
screenshotIfFailed,
8+
DEFAULT_CONNECTION_NAME_1,
9+
} from '../../helpers/compass';
10+
import type { Compass } from '../../helpers/compass';
11+
import * as Selectors from '../../helpers/selectors';
12+
import { createNumbersCollection } from '../../helpers/insert-data';
13+
import { isTestingAtlasCloudSandbox } from '../../helpers/test-runner-context';
14+
import { switchPipelineMode } from '../../helpers/commands/switch-pipeline-mode';
15+
16+
describe('Collection ai query', function () {
17+
let compass: Compass;
18+
let browser: CompassBrowser;
19+
20+
before(function () {
21+
if (!isTestingAtlasCloudSandbox()) {
22+
this.skip();
23+
}
24+
});
25+
26+
afterEach(async function () {
27+
await screenshotIfFailed(compass, this.currentTest);
28+
await cleanup(compass);
29+
});
30+
31+
describe('when the feature is enabled', function () {
32+
beforeEach(async function () {
33+
compass = await init(this.test?.fullTitle());
34+
browser = compass.browser;
35+
await browser.setupDefaultConnections();
36+
37+
await createNumbersCollection();
38+
await browser.connectToDefaults();
39+
await browser.navigateToCollectionTab(
40+
DEFAULT_CONNECTION_NAME_1,
41+
'test',
42+
'numbers',
43+
'Documents'
44+
);
45+
46+
await browser.setFeature('enableGenAIFeaturesAtlasProject', true);
47+
await browser.setFeature(
48+
'enableGenAISampleDocumentPassingOnAtlasProject',
49+
true
50+
);
51+
await browser.setFeature('enableGenAIFeaturesAtlasOrg', true);
52+
await browser.setFeature('optInDataExplorerGenAIFeatures', true);
53+
});
54+
55+
describe('on the documents tab', function () {
56+
beforeEach(async function () {
57+
await browser.navigateToCollectionTab(
58+
DEFAULT_CONNECTION_NAME_1,
59+
'test',
60+
'numbers',
61+
'Documents'
62+
);
63+
});
64+
65+
it('should update the query bar with a generated query', async function () {
66+
// Click the ai entry button.
67+
await browser.clickVisible(Selectors.GenAIEntryButton);
68+
69+
// Enter the ai prompt.
70+
await browser.clickVisible(Selectors.GenAITextInput);
71+
72+
const testUserInput = 'find all documents where i is greater than 50';
73+
await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
74+
75+
// Click generate.
76+
await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
77+
78+
// Wait for the ipc events to succeed.
79+
await browser.waitUntil(async function () {
80+
// Make sure the query bar was updated.
81+
const queryBarFilterContent = await browser.getCodemirrorEditorText(
82+
Selectors.queryBarOptionInputFilter('Documents')
83+
);
84+
return (
85+
queryBarFilterContent.includes('$gt') &&
86+
queryBarFilterContent.includes('50')
87+
);
88+
});
89+
90+
// Run it and check that the correct documents are shown.
91+
await browser.runFind('Documents', true);
92+
const modifiedResult = await browser.getFirstListDocument();
93+
expect(modifiedResult.i).to.be.equal('51');
94+
});
95+
});
96+
97+
describe('on the aggregations tab', function () {
98+
beforeEach(async function () {
99+
await browser.navigateToCollectionTab(
100+
DEFAULT_CONNECTION_NAME_1,
101+
'test',
102+
'numbers',
103+
'Aggregations'
104+
);
105+
106+
await switchPipelineMode(browser, 'as-text');
107+
});
108+
109+
it('should update the aggregation editor with a generated aggregation', async function () {
110+
// Click the ai entry button.
111+
await browser.clickVisible(Selectors.GenAIOpenButton);
112+
113+
// Enter the ai prompt.
114+
await browser.clickVisible(Selectors.GenAITextInput);
115+
116+
const testUserInput = 'find all documents where i is 99';
117+
await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
118+
119+
// Click generate.
120+
await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
121+
122+
await browser.waitUntil(async function () {
123+
const textEditor = browser.$(Selectors.AggregationAsTextEditor);
124+
const textContent = await textEditor.getText();
125+
return textContent.includes('$match');
126+
});
127+
128+
// Run it and check that the correct documents are shown.
129+
await browser.clickVisible(Selectors.RunPipelineButton);
130+
const resultsWorkspace = browser.$(
131+
Selectors.AggregationResultsWorkspace
132+
);
133+
await resultsWorkspace.waitForDisplayed();
134+
135+
await browser.clickVisible(
136+
Selectors.AggregationResultsJSONListSwitchButton
137+
);
138+
const rawDocuments = await browser.getCodemirrorEditorTextAll(
139+
Selectors.DocumentJSONEntry
140+
);
141+
const documents = rawDocuments.map((text) => {
142+
return JSON.parse(text);
143+
});
144+
145+
expect(documents).to.have.lengthOf(1);
146+
expect(documents[0]).to.have.property('_id');
147+
expect(documents[0]).to.have.property('i', 99);
148+
});
149+
});
150+
});
151+
152+
describe('when the org feature is disabled', function () {
153+
beforeEach(async function () {
154+
compass = await init(this.test?.fullTitle());
155+
browser = compass.browser;
156+
await browser.setupDefaultConnections();
157+
158+
await createNumbersCollection();
159+
await browser.connectToDefaults();
160+
await browser.navigateToCollectionTab(
161+
DEFAULT_CONNECTION_NAME_1,
162+
'test',
163+
'numbers',
164+
'Documents'
165+
);
166+
167+
await browser.setFeature('enableGenAIFeaturesAtlasProject', true);
168+
await browser.setFeature(
169+
'enableGenAISampleDocumentPassingOnAtlasProject',
170+
true
171+
);
172+
await browser.setFeature('enableGenAIFeaturesAtlasOrg', false);
173+
await browser.setFeature('optInDataExplorerGenAIFeatures', true);
174+
});
175+
176+
it('should not show the gen ai intro button', async function () {
177+
// Ensure the query bar is shown.
178+
await browser
179+
.$(Selectors.queryBarOptionInputFilter('Documents'))
180+
.waitForDisplayed();
181+
182+
const aiIntroButton = browser.$(Selectors.GenAIEntryButton);
183+
const isSidebarCreateCollectionButtonExisting =
184+
await aiIntroButton.isExisting();
185+
expect(isSidebarCreateCollectionButtonExisting).to.be.equal(false);
186+
});
187+
});
188+
});

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ describe('Atlas Login', function () {
295295
// control over it
296296
await browser.clickVisible('span=Not now');
297297

298-
const aiInput = browser.$(Selectors.QueryBarAITextInput);
298+
const aiInput = browser.$(Selectors.GenAITextInput);
299299
expect(await aiInput.isExisting()).to.eq(false);
300300
expect(await generateQueryButton.isDisplayed()).to.eq(true);
301301
});
@@ -328,7 +328,7 @@ describe('Atlas Login', function () {
328328
// control over it
329329
await browser.clickVisible('span=Not now');
330330

331-
const aiInput = browser.$(Selectors.QueryBarAITextInput);
331+
const aiInput = browser.$(Selectors.GenAITextInput);
332332
expect(await aiInput.isExisting()).to.eq(false);
333333
expect(await generateQueryButton.isDisplayed()).to.eq(true);
334334
});

‎packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
import { saveAggregationPipeline } from '../helpers/commands/save-aggregation-pipeline';
2020
import { Key } from 'webdriverio';
2121
import type { ChainablePromiseElement } from 'webdriverio';
22+
import { switchPipelineMode } from '../helpers/commands/switch-pipeline-mode';
2223

2324
const { expect } = chai;
2425

@@ -83,14 +84,6 @@ async function waitForTab(browser: CompassBrowser, namespace: string) {
8384
);
8485
}
8586

86-
async function switchPipelineMode(
87-
browser: CompassBrowser,
88-
mode: 'as-text' | 'builder-ui'
89-
) {
90-
await browser.clickVisible(Selectors.aggregationPipelineModeToggle(mode));
91-
await browser.waitForAnimations(Selectors.AggregationBuilderWorkspace);
92-
}
93-
9487
async function deleteStage(
9588
browser: CompassBrowser,
9689
index: number

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

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@ describe('Collection ai query', function () {
4545
clearRequests = _clearRequests;
4646
setMockAtlasServerResponse = _setMockAtlasServerResponse;
4747

48-
process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE = endpoint;
4948
process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = endpoint;
5049

5150
telemetry = await startTelemetryServer();
@@ -73,7 +72,6 @@ describe('Collection ai query', function () {
7372

7473
await stopMockAtlasServer();
7574

76-
delete process.env.COMPASS_ATLAS_SERVICE_BASE_URL_OVERRIDE;
7775
delete process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN;
7876

7977
await cleanup(compass);
@@ -101,19 +99,16 @@ describe('Collection ai query', function () {
10199

102100
it('makes request to the server and updates the query bar with the response', async function () {
103101
// Click the ai entry button.
104-
await browser.clickVisible(Selectors.QueryBarAIEntryButton);
102+
await browser.clickVisible(Selectors.GenAIEntryButton);
105103

106104
// Enter the ai prompt.
107-
await browser.clickVisible(Selectors.QueryBarAITextInput);
105+
await browser.clickVisible(Selectors.GenAITextInput);
108106

109107
const testUserInput = 'find all documents where i is greater than 50';
110-
await browser.setValueVisible(
111-
Selectors.QueryBarAITextInput,
112-
testUserInput
113-
);
108+
await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
114109

115110
// Click generate.
116-
await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
111+
await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
117112

118113
// Wait for the ipc events to succeed.
119114
await browser.waitUntil(async function () {
@@ -163,22 +158,19 @@ describe('Collection ai query', function () {
163158

164159
it('the error is shown to the user', async function () {
165160
// Click the ai entry button.
166-
await browser.clickVisible(Selectors.QueryBarAIEntryButton);
161+
await browser.clickVisible(Selectors.GenAIEntryButton);
167162

168163
// Enter the ai prompt.
169-
await browser.clickVisible(Selectors.QueryBarAITextInput);
164+
await browser.clickVisible(Selectors.GenAITextInput);
170165

171166
const testUserInput = 'find all documents where i is greater than 50';
172-
await browser.setValueVisible(
173-
Selectors.QueryBarAITextInput,
174-
testUserInput
175-
);
167+
await browser.setValueVisible(Selectors.GenAITextInput, testUserInput);
176168

177169
// Click generate.
178-
await browser.clickVisible(Selectors.QueryBarAIGenerateQueryButton);
170+
await browser.clickVisible(Selectors.GenAIGenerateQueryButton);
179171

180172
// Check that the error is shown.
181-
const errorBanner = browser.$(Selectors.QueryBarAIErrorMessageBanner);
173+
const errorBanner = browser.$(Selectors.GenAIErrorMessageBanner);
182174
await errorBanner.waitForDisplayed();
183175
expect(await errorBanner.getText()).to.equal(
184176
'Sorry, we were unable to generate the query, please try again. If the error persists, try changing your prompt.'

‎packages/compass-generative-ai/src/components/ai-experience-entry.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ const aiEntryLightModeStyles = css(
8686
);
8787

8888
function AIExperienceEntry({
89-
'data-testid': dataTestId,
89+
'data-testid': dataTestId = 'open-gen-ai-button',
9090
type,
9191
onClick,
9292
}: {

0 commit comments

Comments
 (0)