Skip to content

Commit 8d63966

Browse files
gagikgagik
authored
chore(compass-e2e-tests): add assistant end to end tests COMPASS-9748 (#7429)
1 parent bc8424b commit 8d63966

File tree

11 files changed

+964
-10
lines changed

11 files changed

+964
-10
lines changed

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

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
import http from 'http';
2+
import { once } from 'events';
3+
import type { AddressInfo } from 'net';
4+
5+
export type MockAssistantResponse = {
6+
status: number;
7+
body: string;
8+
};
9+
10+
function sendStreamingResponse(res: http.ServerResponse, content: string) {
11+
// OpenAI Responses API streaming response format using Server-Sent Events
12+
res.writeHead(200, {
13+
'Content-Type': 'text/event-stream; charset=utf-8',
14+
'Cache-Control': 'no-cache',
15+
Connection: 'keep-alive',
16+
'Transfer-Encoding': 'chunked',
17+
});
18+
19+
const responseId = `resp_${Date.now()}`;
20+
const itemId = `item_${Date.now()}`;
21+
let sequenceNumber = 0;
22+
23+
// Send response.created event
24+
res.write(
25+
`data: ${JSON.stringify({
26+
type: 'response.created',
27+
response: {
28+
id: responseId,
29+
object: 'realtime.response',
30+
status: 'in_progress',
31+
output: [],
32+
usage: {
33+
input_tokens: 0,
34+
output_tokens: 0,
35+
total_tokens: 0,
36+
},
37+
},
38+
sequence_number: sequenceNumber++,
39+
})}\n\n`
40+
);
41+
42+
// Send output_item.added event
43+
res.write(
44+
`data: ${JSON.stringify({
45+
type: 'response.output_item.added',
46+
response_id: responseId,
47+
output_index: 0,
48+
item: {
49+
id: itemId,
50+
object: 'realtime.item',
51+
type: 'message',
52+
role: 'assistant',
53+
content: [],
54+
},
55+
sequence_number: sequenceNumber++,
56+
})}\n\n`
57+
);
58+
59+
// Send the content in chunks
60+
const words = content.split(' ');
61+
let index = 0;
62+
63+
const sendChunk = () => {
64+
if (index < words.length) {
65+
const word = words[index] + (index < words.length - 1 ? ' ' : '');
66+
// Send output_text.delta event
67+
res.write(
68+
`data: ${JSON.stringify({
69+
type: 'response.output_text.delta',
70+
response_id: responseId,
71+
item_id: itemId,
72+
output_index: 0,
73+
delta: word,
74+
sequence_number: sequenceNumber++,
75+
})}\n\n`
76+
);
77+
index++;
78+
setTimeout(sendChunk, 10);
79+
} else {
80+
// Send output_item.done event
81+
res.write(
82+
`data: ${JSON.stringify({
83+
type: 'response.output_item.done',
84+
response_id: responseId,
85+
output_index: 0,
86+
item: {
87+
id: itemId,
88+
object: 'realtime.item',
89+
type: 'message',
90+
role: 'assistant',
91+
content: [
92+
{
93+
type: 'text',
94+
text: content,
95+
},
96+
],
97+
},
98+
sequence_number: sequenceNumber++,
99+
})}\n\n`
100+
);
101+
102+
// Send response.completed event
103+
const tokenCount = Math.ceil(content.split(' ').length * 1.3);
104+
res.write(
105+
`data: ${JSON.stringify({
106+
type: 'response.completed',
107+
response: {
108+
id: responseId,
109+
object: 'realtime.response',
110+
status: 'completed',
111+
output: [
112+
{
113+
id: itemId,
114+
object: 'realtime.item',
115+
type: 'message',
116+
role: 'assistant',
117+
content: [
118+
{
119+
type: 'text',
120+
text: content,
121+
},
122+
],
123+
},
124+
],
125+
usage: {
126+
input_tokens: 10,
127+
output_tokens: tokenCount,
128+
total_tokens: 10 + tokenCount,
129+
},
130+
},
131+
sequence_number: sequenceNumber++,
132+
})}\n\n`
133+
);
134+
135+
res.write('data: [DONE]\n\n');
136+
res.end();
137+
}
138+
};
139+
140+
sendChunk();
141+
}
142+
143+
export const MOCK_ASSISTANT_SERVER_PORT = 27097;
144+
export async function startMockAssistantServer(
145+
{
146+
response: _response,
147+
}: {
148+
response: MockAssistantResponse;
149+
} = {
150+
response: {
151+
status: 200,
152+
body: 'This is a test response from the AI assistant.',
153+
},
154+
}
155+
): Promise<{
156+
clearRequests: () => void;
157+
getResponse: () => MockAssistantResponse;
158+
setResponse: (response: MockAssistantResponse) => void;
159+
getRequests: () => {
160+
content: any;
161+
req: any;
162+
}[];
163+
endpoint: string;
164+
server: http.Server;
165+
stop: () => Promise<void>;
166+
}> {
167+
let requests: {
168+
content: any;
169+
req: any;
170+
}[] = [];
171+
let response = _response;
172+
const server = http
173+
.createServer((req, res) => {
174+
res.setHeader('Access-Control-Allow-Origin', '*');
175+
res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS');
176+
res.setHeader(
177+
'Access-Control-Allow-Headers',
178+
'Content-Type, Authorization, X-Request-Origin, User-Agent'
179+
);
180+
181+
// Handle preflight requests
182+
if (req.method === 'OPTIONS') {
183+
res.writeHead(200);
184+
res.end();
185+
return;
186+
}
187+
188+
// Only handle POST requests for chat completions
189+
if (req.method !== 'POST') {
190+
res.writeHead(404);
191+
return res.end('Not Found');
192+
}
193+
194+
let body = '';
195+
req
196+
.setEncoding('utf8')
197+
.on('data', (chunk) => {
198+
body += chunk;
199+
})
200+
.on('end', () => {
201+
let jsonObject;
202+
try {
203+
jsonObject = JSON.parse(body);
204+
} catch {
205+
res.writeHead(400);
206+
res.setHeader('Content-Type', 'application/json');
207+
return res.end(JSON.stringify({ error: 'Invalid JSON' }));
208+
}
209+
210+
requests.push({
211+
req,
212+
content: jsonObject,
213+
});
214+
215+
if (response.status !== 200) {
216+
res.writeHead(response.status);
217+
res.setHeader('Content-Type', 'application/json');
218+
return res.end(JSON.stringify({ error: response.body }));
219+
}
220+
221+
// Send streaming response
222+
return sendStreamingResponse(res, response.body);
223+
});
224+
})
225+
.listen(MOCK_ASSISTANT_SERVER_PORT);
226+
await once(server, 'listening');
227+
228+
// address() returns either a string or AddressInfo.
229+
const address = server.address() as AddressInfo;
230+
231+
const endpoint = `http://localhost:${address.port}`;
232+
233+
async function stop() {
234+
server.close();
235+
await once(server, 'close');
236+
}
237+
238+
function clearRequests() {
239+
requests = [];
240+
}
241+
242+
function getRequests() {
243+
return requests;
244+
}
245+
246+
function getResponse() {
247+
return response;
248+
}
249+
250+
function setResponse(newResponse: MockAssistantResponse) {
251+
response = newResponse;
252+
}
253+
254+
return {
255+
clearRequests,
256+
getRequests,
257+
endpoint,
258+
server,
259+
getResponse,
260+
setResponse,
261+
stop,
262+
};
263+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ export * from './scroll-to-virtual-item';
3434
export * from './set-export-filename';
3535
export * from './get-feature';
3636
export * from './set-feature';
37+
export * from './set-env';
3738
export * from './save-favorite';
3839
export * from './save-connection-string-as-favorite';
3940
export * from './sidebar-connection';

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import type { CompassBrowser } from '../compass-browser';
2+
import { isTestingWeb } from '../test-runner-context';
3+
4+
/**
5+
* Sets an environment variable override in Compass Web.
6+
* This is only supported in Compass Web tests, not in Compass Desktop.
7+
*
8+
* @example
9+
* // Set the Atlas service URL override in a test
10+
* await browser.setEnv(
11+
* 'COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE',
12+
* mockAtlasServer.endpoint
13+
* );
14+
*
15+
* @param browser - The CompassBrowser instance
16+
* @param key - The environment variable name
17+
* @param value - The environment variable value
18+
*/
19+
export async function setEnv(
20+
browser: CompassBrowser,
21+
key: string,
22+
value: string
23+
): Promise<void> {
24+
if (isTestingWeb()) {
25+
// When running in Compass web we use a global function to set env vars
26+
await browser.execute(
27+
(_key, _value) => {
28+
const kSandboxSetEnvFn = Symbol.for('@compass-web-sandbox-set-env');
29+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
30+
(globalThis as any)[kSandboxSetEnvFn]?.(_key, _value);
31+
},
32+
key,
33+
value
34+
);
35+
return;
36+
}
37+
38+
// When running in Compass desktop, we can't dynamically change env vars
39+
// after the process has started, so we throw an error
40+
throw new Error(
41+
'setEnv is only supported in Compass web. For Compass desktop, set environment variables before starting the app.'
42+
);
43+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
} from './test-runner-paths';
1212
import type { ConnectionInfo } from '@mongodb-js/connection-info';
1313
import ConnectionString from 'mongodb-connection-string-url';
14+
import { MOCK_ASSISTANT_SERVER_PORT } from './assistant-service';
1415

1516
const debug = Debug('compass-e2e-tests:compass-web-sandbox');
1617

@@ -21,6 +22,9 @@ const debug = Debug('compass-e2e-tests:compass-web-sandbox');
2122
process.env.OPEN_BROWSER = 'false'; // tell webpack dev server not to open the default browser
2223
process.env.DISABLE_DEVSERVER_OVERLAY = 'true';
2324
process.env.APP_ENV = 'webdriverio';
25+
// Set the assistant base URL override for tests so we can mock the assistant server
26+
process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE = `http://localhost:${MOCK_ASSISTANT_SERVER_PORT}`;
27+
process.env.COMPASS_OVERRIDE_ENABLE_AI_FEATURES = 'true';
2428

2529
const wait = (ms: number) => {
2630
return new Promise((resolve) => {

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ export const SettingsModalTabSelector = (name: string) =>
1616
`${SettingsModal} [data-testid="sidebar-${name}-item"]`;
1717
export const GeneralSettingsButton = SettingsModalTabSelector('general');
1818
export const GeneralSettingsContent = `${SettingsModal} [data-testid="general-settings"]`;
19+
export const ArtificialIntelligenceSettingsButton =
20+
SettingsModalTabSelector('ai');
21+
export const ArtificialIntelligenceSettingsContent = `${SettingsModal} [data-testid="gen-ai-settings"]`;
1922

2023
export const SettingsInputElement = (settingName: string): string => {
2124
return `${SettingsModal} [data-testid="${settingName}"]`;
@@ -894,6 +897,9 @@ export const AggregationSavedPipelineCardDeleteButton = (
894897
export const AggregationExplainButton =
895898
'[data-testid="pipeline-toolbar-explain-aggregation-button"]';
896899
export const AggregationExplainModal = '[data-testid="explain-plan-modal"]';
900+
export const ExplainPlanInterpretButton =
901+
'[data-testid="interpret-for-me-button"]';
902+
export const ExplainPlanCloseButton = '[data-testid="explain-close-button"]';
897903
export const AggregationExplainModalCloseButton = `${AggregationExplainModal} [aria-label*="Close"]`;
898904

899905
// Create view from pipeline modal
@@ -1510,8 +1516,19 @@ export const SideDrawerCloseButton = `[data-testid="${
15101516
}"]`;
15111517

15121518
// Assistant
1519+
export const AssistantDrawerButton = 'button[aria-label="MongoDB Assistant"]';
1520+
export const AssistantDrawerCloseButton = `[data-testid="lg-drawer-close_button"]`;
15131521
export const AssistantChatMessages = '[data-testid="assistant-chat-messages"]';
1522+
export const AssistantChatMessage = '[data-testid^="assistant-message-"]';
1523+
export const AssistantChatInput = '[data-testid="assistant-chat-input"]';
1524+
export const AssistantChatInputTextArea = `${AssistantChatInput} textarea`;
1525+
export const AssistantChatSubmitButton = `${AssistantChatInput} button[aria-label="Send message"]`;
15141526
export const AssistantClearChatButton = '[data-testid="assistant-clear-chat"]';
1515-
export const ConfirmClearChatModal =
1527+
export const AssistantConfirmClearChatModal =
15161528
'[data-testid="assistant-confirm-clear-chat-modal"]';
1517-
export const ConfirmClearChatModalConfirmButton = `${ConfirmClearChatModal} [data-testid="lg-confirmation_modal-footer-confirm_button"]`;
1529+
export const AssistantConfirmClearChatModalConfirmButton = `${AssistantConfirmClearChatModal} [data-testid="lg-confirmation_modal-footer-confirm_button"]`;
1530+
1531+
// AI Opt-in Modal
1532+
export const AIOptInModal = '[data-testid="ai-optin-modal"]';
1533+
export const AIOptInModalAcceptButton = 'button=Use AI Features';
1534+
export const AIOptInModalDeclineLink = 'span=Not now';

0 commit comments

Comments
 (0)