Skip to content

Commit a656e36

Browse files
gagikgagik
committed
chore(compass-e2e-tests): add assistant end to end tests COMPASS-9384 
1 parent 6147765 commit a656e36

File tree

5 files changed

+666
-3
lines changed

5 files changed

+666
-3
lines changed

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

Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
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 async function startMockAssistantServer(
144+
{
145+
response: _response,
146+
}: {
147+
response: MockAssistantResponse;
148+
} = {
149+
response: {
150+
status: 200,
151+
body: 'This is a test response from the AI assistant.',
152+
},
153+
}
154+
): Promise<{
155+
clearRequests: () => void;
156+
getResponse: () => MockAssistantResponse;
157+
setResponse: (response: MockAssistantResponse) => void;
158+
getRequests: () => {
159+
content: any;
160+
req: any;
161+
}[];
162+
endpoint: string;
163+
server: http.Server;
164+
stop: () => Promise<void>;
165+
}> {
166+
let requests: {
167+
content: any;
168+
req: any;
169+
}[] = [];
170+
let response = _response;
171+
const server = http
172+
.createServer((req, res) => {
173+
// Only handle POST requests for chat completions
174+
if (req.method !== 'POST') {
175+
res.writeHead(404);
176+
return res.end('Not Found');
177+
}
178+
179+
let body = '';
180+
req
181+
.setEncoding('utf8')
182+
.on('data', (chunk) => {
183+
body += chunk;
184+
})
185+
.on('end', () => {
186+
let jsonObject;
187+
try {
188+
jsonObject = JSON.parse(body);
189+
} catch {
190+
res.writeHead(400);
191+
res.setHeader('Content-Type', 'application/json');
192+
return res.end(JSON.stringify({ error: 'Invalid JSON' }));
193+
}
194+
195+
requests.push({
196+
req,
197+
content: jsonObject,
198+
});
199+
200+
if (response.status !== 200) {
201+
res.writeHead(response.status);
202+
res.setHeader('Content-Type', 'application/json');
203+
return res.end(JSON.stringify({ error: response.body }));
204+
}
205+
206+
// Send streaming response
207+
return sendStreamingResponse(res, response.body);
208+
});
209+
})
210+
.listen(0);
211+
await once(server, 'listening');
212+
213+
// address() returns either a string or AddressInfo.
214+
const address = server.address() as AddressInfo;
215+
216+
const endpoint = `http://localhost:${address.port}`;
217+
218+
async function stop() {
219+
server.close();
220+
await once(server, 'close');
221+
}
222+
223+
function clearRequests() {
224+
requests = [];
225+
}
226+
227+
function getRequests() {
228+
return requests;
229+
}
230+
231+
function getResponse() {
232+
return response;
233+
}
234+
235+
function setResponse(newResponse: MockAssistantResponse) {
236+
response = newResponse;
237+
}
238+
239+
return {
240+
clearRequests,
241+
getRequests,
242+
endpoint,
243+
server,
244+
getResponse,
245+
setResponse,
246+
stop,
247+
};
248+
}

‎packages/compass-e2e-tests/helpers/commands/open-settings-modal.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1+
import type { ChainablePromiseElement } from 'webdriverio';
12
import type { CompassBrowser } from '../compass-browser';
23
import * as Selectors from '../selectors';
34

45
export async function openSettingsModal(
56
browser: CompassBrowser,
67
tab?: string
7-
): Promise<void> {
8+
): Promise<ChainablePromiseElement> {
89
await browser.execute(() => {
910
// eslint-disable-next-line @typescript-eslint/no-require-imports
1011
require('electron').ipcRenderer.emit('window:show-settings');
@@ -15,4 +16,5 @@ export async function openSettingsModal(
1516
if (tab) {
1617
await browser.clickVisible(Selectors.SettingsModalTabSelector(tab));
1718
}
19+
return settingsModalElement;
1820
}

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

Lines changed: 16 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}"]`;
@@ -1510,8 +1513,19 @@ export const SideDrawerCloseButton = `[data-testid="${
15101513
}"]`;
15111514

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

0 commit comments

Comments
 (0)