Skip to content

Commit 887bec6

Browse files
committed
test: add tests for chat route
Signed-off-by: tygao <[email protected]>
1 parent 7b8ec87 commit 887bec6

File tree

1 file changed

+174
-11
lines changed

1 file changed

+174
-11
lines changed

server/routes/chat_routes.test.ts

+174-11
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,24 @@ import { mockOllyChatService } from '../services/chat/olly_chat_service.mock';
1414
import { loggerMock } from '../../../../src/core/server/logging/logger.mock';
1515
import { registerChatRoutes } from './chat_routes';
1616
import { ASSISTANT_API } from '../../common/constants/llm';
17+
import { getOpenSearchClientTransport } from '../utils/get_opensearch_client_transport';
1718

18-
const mockedLogger = loggerMock.create();
19+
jest.mock('../utils/get_opensearch_client_transport');
20+
21+
beforeEach(() => {
22+
(getOpenSearchClientTransport as jest.Mock).mockImplementation(({ dataSourceId }) => {
23+
if (dataSourceId) {
24+
return 'dataSource-client';
25+
} else {
26+
return 'client';
27+
}
28+
});
29+
});
30+
afterEach(() => {
31+
(getOpenSearchClientTransport as jest.Mock).mockClear();
32+
});
1933

34+
const mockedLogger = loggerMock.create();
2035
const router = new Router(
2136
'',
2237
mockedLogger,
@@ -30,38 +45,90 @@ registerChatRoutes(router, {
3045
messageParsers: [],
3146
});
3247

33-
const triggerDeleteConversation = (conversationId: string) =>
48+
const triggerDeleteConversation = (conversationId: string, dataSourceId?: string) =>
3449
triggerHandler(router, {
3550
method: 'delete',
3651
path: `${ASSISTANT_API.CONVERSATION}/{conversationId}`,
37-
req: httpServerMock.createRawRequest({ params: { conversationId } }),
52+
req: httpServerMock.createRawRequest({
53+
params: { conversationId },
54+
...(dataSourceId
55+
? {
56+
query: {
57+
dataSourceId,
58+
},
59+
}
60+
: {}),
61+
}),
3862
});
3963
const triggerUpdateConversation = (
4064
params: { conversationId: string },
41-
payload: { title: string }
65+
payload: { title: string },
66+
dataSourceId?: string
4267
) =>
4368
triggerHandler(router, {
4469
method: 'put',
4570
path: `${ASSISTANT_API.CONVERSATION}/{conversationId}`,
46-
req: httpServerMock.createRawRequest({ params, payload }),
71+
req: httpServerMock.createRawRequest({
72+
params,
73+
payload,
74+
...(dataSourceId
75+
? {
76+
query: {
77+
dataSourceId,
78+
},
79+
}
80+
: {}),
81+
}),
4782
});
48-
const triggerGetTrace = (interactionId: string) =>
83+
const triggerGetTrace = (interactionId: string, dataSourceId?: string) =>
4984
triggerHandler(router, {
5085
method: 'get',
5186
path: `${ASSISTANT_API.TRACE}/{interactionId}`,
52-
req: httpServerMock.createRawRequest({ params: { interactionId } }),
87+
req: httpServerMock.createRawRequest({
88+
params: { interactionId },
89+
...(dataSourceId
90+
? {
91+
query: {
92+
dataSourceId,
93+
},
94+
}
95+
: {}),
96+
}),
5397
});
54-
const triggerAbortAgentExecution = (conversationId: string) =>
98+
const triggerAbortAgentExecution = (conversationId: string, dataSourceId?: string) =>
5599
triggerHandler(router, {
56100
method: 'post',
57101
path: ASSISTANT_API.ABORT_AGENT_EXECUTION,
58-
req: httpServerMock.createRawRequest({ payload: { conversationId } }),
102+
req: httpServerMock.createRawRequest({
103+
payload: { conversationId },
104+
...(dataSourceId
105+
? {
106+
query: {
107+
dataSourceId,
108+
},
109+
}
110+
: {}),
111+
}),
59112
});
60-
const triggerFeedback = (params: { interactionId: string }, payload: { satisfaction: boolean }) =>
113+
const triggerFeedback = (
114+
params: { interactionId: string },
115+
payload: { satisfaction: boolean },
116+
dataSourceId?: string
117+
) =>
61118
triggerHandler(router, {
62119
method: 'put',
63120
path: `${ASSISTANT_API.FEEDBACK}/{interactionId}`,
64-
req: httpServerMock.createRawRequest({ params, payload }),
121+
req: httpServerMock.createRawRequest({
122+
params,
123+
payload,
124+
...(dataSourceId
125+
? {
126+
query: {
127+
dataSourceId,
128+
},
129+
}
130+
: {}),
131+
}),
65132
});
66133

67134
describe('chat routes', () => {
@@ -80,6 +147,22 @@ describe('chat routes', () => {
80147

81148
expect(mockAgentFrameworkStorageService.deleteConversation).not.toHaveBeenCalled();
82149
const result = (await triggerDeleteConversation('foo')) as ResponseObject;
150+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('client');
151+
expect(mockAgentFrameworkStorageService.deleteConversation).toHaveBeenCalledWith('foo');
152+
expect(result.source).toMatchInlineSnapshot(`
153+
Object {
154+
"success": true,
155+
}
156+
`);
157+
});
158+
159+
it('should call delete conversation with passed data source id and get data source transport', async () => {
160+
mockAgentFrameworkStorageService.deleteConversation.mockResolvedValueOnce({
161+
success: true,
162+
});
163+
expect(mockAgentFrameworkStorageService.deleteConversation).not.toHaveBeenCalled();
164+
const result = (await triggerDeleteConversation('foo', 'data_source_id')) as ResponseObject;
165+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('dataSource-client');
83166
expect(mockAgentFrameworkStorageService.deleteConversation).toHaveBeenCalledWith('foo');
84167
expect(result.source).toMatchInlineSnapshot(`
85168
Object {
@@ -109,6 +192,7 @@ describe('chat routes', () => {
109192
{ conversationId: 'foo' },
110193
{ title: 'new-title' }
111194
)) as ResponseObject;
195+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('client');
112196
expect(mockAgentFrameworkStorageService.updateConversation).toHaveBeenCalledWith(
113197
'foo',
114198
'new-title'
@@ -120,6 +204,29 @@ describe('chat routes', () => {
120204
`);
121205
});
122206

207+
it('should call update conversation with passed data source id and title then get data source transport', async () => {
208+
mockAgentFrameworkStorageService.updateConversation.mockResolvedValueOnce({
209+
success: true,
210+
});
211+
212+
expect(mockAgentFrameworkStorageService.updateConversation).not.toHaveBeenCalled();
213+
const result = (await triggerUpdateConversation(
214+
{ conversationId: 'foo' },
215+
{ title: 'new-title' },
216+
'data_source_id'
217+
)) as ResponseObject;
218+
expect(mockAgentFrameworkStorageService.updateConversation).toHaveBeenCalledWith(
219+
'foo',
220+
'new-title'
221+
);
222+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('dataSource-client');
223+
expect(result.source).toMatchInlineSnapshot(`
224+
Object {
225+
"success": true,
226+
}
227+
`);
228+
});
229+
123230
it('should log error and return 500 error when failed to update conversation', async () => {
124231
mockAgentFrameworkStorageService.updateConversation.mockRejectedValueOnce(new Error());
125232

@@ -149,6 +256,27 @@ describe('chat routes', () => {
149256

150257
expect(mockAgentFrameworkStorageService.getTraces).not.toHaveBeenCalled();
151258
const result = (await triggerGetTrace('interaction-1')) as ResponseObject;
259+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('client');
260+
expect(mockAgentFrameworkStorageService.getTraces).toHaveBeenCalledWith('interaction-1');
261+
expect(result.source).toEqual(getTraceResultMock);
262+
});
263+
264+
it('should call get traces with passed data source id and get data source transport', async () => {
265+
const getTraceResultMock = [
266+
{
267+
interactionId: 'interaction-1',
268+
createTime: '',
269+
input: 'foo',
270+
output: 'bar',
271+
origin: '',
272+
traceNumber: 0,
273+
},
274+
];
275+
mockAgentFrameworkStorageService.getTraces.mockResolvedValueOnce(getTraceResultMock);
276+
277+
expect(mockAgentFrameworkStorageService.getTraces).not.toHaveBeenCalled();
278+
const result = (await triggerGetTrace('interaction-1', 'data_source_id')) as ResponseObject;
279+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('dataSource-client');
152280
expect(mockAgentFrameworkStorageService.getTraces).toHaveBeenCalledWith('interaction-1');
153281
expect(result.source).toEqual(getTraceResultMock);
154282
});
@@ -169,6 +297,16 @@ describe('chat routes', () => {
169297

170298
await triggerAbortAgentExecution('foo');
171299
expect(mockOllyChatService.abortAgentExecution).toHaveBeenCalledWith('foo');
300+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('client');
301+
expect(mockedLogger.info).toHaveBeenCalledWith('Abort agent execution: foo');
302+
});
303+
304+
it('should call get abort agent with passed data source id and get data source transport ', async () => {
305+
expect(mockOllyChatService.abortAgentExecution).not.toHaveBeenCalled();
306+
307+
await triggerAbortAgentExecution('foo', 'data_source_id');
308+
expect(mockOllyChatService.abortAgentExecution).toHaveBeenCalledWith('foo');
309+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('dataSource-client');
172310
expect(mockedLogger.info).toHaveBeenCalledWith('Abort agent execution: foo');
173311
});
174312

@@ -200,6 +338,31 @@ describe('chat routes', () => {
200338
{ interactionId: 'foo' },
201339
{ satisfaction: true }
202340
)) as ResponseObject;
341+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('client');
342+
expect(mockAgentFrameworkStorageService.updateInteraction).toHaveBeenCalledWith('foo', {
343+
feedback: {
344+
satisfaction: true,
345+
},
346+
});
347+
expect(result.source).toMatchInlineSnapshot(`
348+
Object {
349+
"success": true,
350+
}
351+
`);
352+
});
353+
354+
it('should call update interaction with passed data source id and get data source transport', async () => {
355+
mockAgentFrameworkStorageService.updateConversation.mockResolvedValueOnce({
356+
success: true,
357+
});
358+
359+
expect(mockAgentFrameworkStorageService.updateConversation).not.toHaveBeenCalled();
360+
const result = (await triggerFeedback(
361+
{ interactionId: 'foo' },
362+
{ satisfaction: true },
363+
'data_source_id'
364+
)) as ResponseObject;
365+
expect(getOpenSearchClientTransport.mock.results[0].value).toBe('dataSource-client');
203366
expect(mockAgentFrameworkStorageService.updateInteraction).toHaveBeenCalledWith('foo', {
204367
feedback: {
205368
satisfaction: true,

0 commit comments

Comments
 (0)