lobe-chat

1
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
2

3
import { ChatMessage } from '@/types/message';
4

5
import { CreateMessageParams, MessageModel } from '../message';
6

7
describe('MessageModel', () => {
8
  let messageData: CreateMessageParams;
9

10
  beforeEach(() => {
11
    // 设置正确结构的消息数据
12
    messageData = {
13
      content: 'Test message content',
14
      role: 'user',
15
      sessionId: 'session1',
16
      topicId: 'topic1',
17
    };
18
  });
19

20
  afterEach(async () => {
21
    // 每次测试后清理数据库
22
    await MessageModel.clearTable();
23
  });
24

25
  describe('create', () => {
26
    it('should create a message record', async () => {
27
      const result = await MessageModel.create(messageData);
28

29
      expect(result).toHaveProperty('id');
30
      // 验证消息是否已添加到数据库
31
      const messageInDb = await MessageModel.findById(result.id);
32

33
      expect(messageInDb).toEqual(
34
        expect.objectContaining({
35
          content: messageData.content,
36
          role: messageData.role,
37
          sessionId: messageData.sessionId,
38
          topicId: messageData.topicId,
39
        }),
40
      );
41
    });
42

43
    it('should create with tts', async () => {
44
      const result = await MessageModel.create({
45
        content: 'abc',
46
        role: 'assistant',
47
        extra: { translate: { content: 'avc', from: 'a', to: 'f' } },
48
        sessionId: 'a',
49
      });
50

51
      // 验证消息是否已添加到数据库
52
      const messageInDb = await MessageModel.findById(result.id);
53

54
      expect(messageInDb).toEqual(
55
        expect.objectContaining({
56
          content: 'abc',
57
          role: 'assistant',
58
          translate: { content: 'avc', from: 'a', to: 'f' },
59
          sessionId: 'a',
60
        }),
61
      );
62
    });
63
  });
64

65
  describe('batchCreate', () => {
66
    it('should batch create message records', async () => {
67
      const messagesToCreate = [messageData, messageData] as ChatMessage[];
68
      const results = await MessageModel.batchCreate(messagesToCreate);
69

70
      expect(results.success).toBeTruthy();
71
      expect(results.errors).toBeUndefined();
72

73
      // 验证消息是否已添加到数据库
74
      for (const message of results.ids!) {
75
        const messageInDb = await MessageModel.findById(message);
76
        expect(messageInDb).toEqual(
77
          expect.objectContaining({
78
            content: messageData.content,
79
            role: messageData.role,
80
            sessionId: messageData.sessionId,
81
            topicId: messageData.topicId,
82
          }),
83
        );
84
      }
85
    });
86
  });
87

88
  describe('query', () => {
89
    it('should query messages with pagination', async () => {
90
      // 创建多条消息以测试查询方法
91
      await MessageModel.batchCreate([messageData, messageData] as ChatMessage[]);
92

93
      const queriedMessages = await MessageModel.query({
94
        pageSize: 1,
95
        current: 0,
96
        sessionId: messageData.sessionId,
97
        topicId: messageData.topicId,
98
      });
99

100
      expect(queriedMessages).toHaveLength(1);
101
    });
102

103
    it('should query correctly without topic id', async () => {
104
      // 创建多条消息以测试查询方法
105
      await MessageModel.batchCreate([messageData, messageData] as ChatMessage[]);
106

107
      const queriedMessages = await MessageModel.query({ sessionId: messageData.sessionId });
108

109
      expect(queriedMessages).toHaveLength(0);
110
    });
111

112
    it('should query correctly with exactly topic id', async () => {
113
      // 创建多条消息以测试查询方法
114
      await MessageModel.batchCreate([
115
        messageData,
116
        { ...messageData, topicId: undefined },
117
      ] as ChatMessage[]);
118

119
      const queriedMessages = await MessageModel.query({ sessionId: messageData.sessionId });
120

121
      expect(queriedMessages).toHaveLength(1);
122
    });
123

124
    it('should should have correct order', async () => {
125
      const data: ChatMessage[] = [
126
        {
127
          role: 'user',
128
          content: '1',
129
          createdAt: 1697120044345,
130
          id: 'NQ7RscYx',
131
          updatedAt: 1697120181827,
132
          extra: {},
133
          meta: {},
134
          sessionId: '1',
135
        },
136
        {
137
          role: 'assistant',
138
          content: '2',
139
          parentId: 'NQ7RscYx',
140
          createdAt: 1697120130973,
141
          id: '9tDAumEx',
142
          updatedAt: 1697120181827,
143
          meta: {},
144
          extra: {
145
            fromModel: 'gpt-3.5-turbo-16k',
146
          },
147
          sessionId: '1',
148
        },
149
        {
150
          role: 'assistant',
151
          content: '3',
152
          parentId: 'tOMH7c5R',
153
          meta: {},
154
          createdAt: 1697120163272,
155
          id: '5Ie5hClg',
156
          updatedAt: 1697120181827,
157
          extra: {
158
            fromModel: 'gpt-3.5-turbo-16k',
159
          },
160
          sessionId: '1',
161
        },
162
        {
163
          role: 'user',
164
          content: '4',
165
          meta: {},
166
          createdAt: 1697120163272,
167
          id: 'tOMH7c5R',
168
          updatedAt: 1697120181827,
169
          extra: {},
170
          sessionId: '1',
171
        },
172
      ];
173

174
      await MessageModel.batchCreate(data);
175

176
      const queriedMessages = await MessageModel.query({ sessionId: '1' });
177

178
      expect(queriedMessages).toEqual([
179
        {
180
          role: 'user',
181
          content: '1',
182
          createdAt: 1697120044345,
183
          id: 'NQ7RscYx',
184
          updatedAt: 1697120181827,
185
          sessionId: '1',
186
          extra: {},
187
          meta: {},
188
        },
189
        {
190
          role: 'assistant',
191
          content: '2',
192
          parentId: 'NQ7RscYx',
193
          createdAt: 1697120130973,
194
          id: '9tDAumEx',
195
          sessionId: '1',
196
          updatedAt: 1697120181827,
197
          meta: {},
198
          extra: {
199
            fromModel: 'gpt-3.5-turbo-16k',
200
          },
201
        },
202
        {
203
          role: 'user',
204
          content: '4',
205
          sessionId: '1',
206
          createdAt: 1697120163272,
207
          id: 'tOMH7c5R',
208
          updatedAt: 1697120181827,
209
          meta: {},
210
          extra: {},
211
        },
212
        {
213
          role: 'assistant',
214
          content: '3',
215
          parentId: 'tOMH7c5R',
216
          meta: {},
217
          createdAt: 1697120163272,
218
          sessionId: '1',
219
          id: '5Ie5hClg',
220
          updatedAt: 1697120181827,
221
          extra: {
222
            fromModel: 'gpt-3.5-turbo-16k',
223
          },
224
        },
225
      ]);
226
    });
227
  });
228

229
  describe('findById', () => {
230
    it('should find a message by id', async () => {
231
      const createdMessage = await MessageModel.create(messageData);
232
      const messageInDb = await MessageModel.findById(createdMessage.id);
233

234
      expect(messageInDb).toEqual(
235
        expect.objectContaining({
236
          id: createdMessage.id,
237
          content: messageData.content,
238
        }),
239
      );
240
    });
241
  });
242

243
  describe('delete', () => {
244
    it('should delete a message', async () => {
245
      const createdMessage = await MessageModel.create(messageData);
246
      await MessageModel.delete(createdMessage.id);
247

248
      const messageInDb = await MessageModel.findById(createdMessage.id);
249
      expect(messageInDb).toBeUndefined();
250
    });
251
  });
252

253
  describe('bulkDelete', () => {
254
    it('should delete many messages', async () => {
255
      const createdMessage = await MessageModel.create(messageData);
256
      const createdMessage2 = await MessageModel.create(messageData);
257
      await MessageModel.bulkDelete([createdMessage.id, createdMessage2.id]);
258

259
      const messageInDb1 = await MessageModel.findById(createdMessage.id);
260
      const messageInDb2 = await MessageModel.findById(createdMessage2.id);
261
      expect(messageInDb1).toBeUndefined();
262
      expect(messageInDb2).toBeUndefined();
263
    });
264
  });
265

266
  describe('update', () => {
267
    it('should update a message', async () => {
268
      const createdMessage = await MessageModel.create(messageData);
269
      const updateData = { content: 'Updated content' };
270

271
      await MessageModel.update(createdMessage.id, updateData);
272
      const updatedMessage = await MessageModel.findById(createdMessage.id);
273

274
      expect(updatedMessage).toHaveProperty('content', 'Updated content');
275
    });
276

277
    it('should update a role and plugins', async () => {
278
      const createdMessage = await MessageModel.create(messageData);
279
      const updateData = {
280
        role: 'tool' as const,
281
        plugin: { apiName: 'a', identifier: 'b', arguments: 'abc' },
282
      };
283

284
      await MessageModel.update(createdMessage.id, updateData);
285
      const updatedMessage = await MessageModel.findById(createdMessage.id);
286

287
      expect(updatedMessage).toHaveProperty('role', 'tool');
288
    });
289
  });
290

291
  describe('batchUpdate', () => {
292
    it('should batch update messages', async () => {
293
      const createdMessage1 = await MessageModel.create(messageData);
294
      const createdMessage2 = await MessageModel.create(messageData);
295
      const updateData = { content: 'Batch updated content' };
296

297
      const numUpdated = await MessageModel.batchUpdate(
298
        [createdMessage1.id, createdMessage2.id],
299
        updateData,
300
      );
301

302
      expect(numUpdated).toBe(2);
303

304
      const updatedMessage1 = await MessageModel.findById(createdMessage1.id);
305
      const updatedMessage2 = await MessageModel.findById(createdMessage2.id);
306

307
      expect(updatedMessage1).toHaveProperty('content', 'Batch updated content');
308
      expect(updatedMessage2).toHaveProperty('content', 'Batch updated content');
309
    });
310
  });
311

312
  describe('batchDelete', () => {
313
    it('should batch delete messages by session id', async () => {
314
      // 创建多条消息以测试批量删除方法
315
      const createdMessage1 = await MessageModel.create(messageData);
316
      const createdMessage2 = await MessageModel.create(messageData);
317

318
      await MessageModel.batchDelete(messageData.sessionId, undefined);
319

320
      // 验证所有具有给定会话 ID 的消息是否已删除
321
      const messagesInDb = await MessageModel.query({ sessionId: messageData.sessionId });
322
      expect(messagesInDb).toHaveLength(0);
323
    });
324

325
    it('should batch delete messages by session id and topic id', async () => {
326
      // 创建多条消息以测试批量删除方法
327
      const createdMessage1 = await MessageModel.create(messageData);
328
      const createdMessage2 = await MessageModel.create(messageData);
329

330
      await MessageModel.batchDelete(messageData.sessionId, messageData.topicId);
331

332
      // 验证所有具有给定会话 ID 和话题 ID 的消息是否已删除
333
      const messagesInDb = await MessageModel.query({
334
        sessionId: messageData.sessionId,
335
        topicId: messageData.topicId,
336
      });
337
      expect(messagesInDb).toHaveLength(0);
338
    });
339
  });
340

341
  describe('duplicateMessages', () => {
342
    it('should duplicate messages and update parentId for copied messages', async () => {
343
      // 创建原始消息和父消息
344
      const parentMessageData: CreateMessageParams = {
345
        content: 'Parent message content',
346
        role: 'user',
347
        sessionId: 'session1',
348
        topicId: undefined,
349
      };
350
      const parentMessage = await MessageModel.create(parentMessageData);
351

352
      const childMessageData: CreateMessageParams = {
353
        content: 'Child message content',
354
        role: 'user',
355
        sessionId: 'session1',
356
        parentId: parentMessage.id,
357
      };
358

359
      await MessageModel.create(childMessageData);
360

361
      // 获取数据库中的消息以进行复制
362
      const originalMessages = await MessageModel.queryAll();
363

364
      // 执行复制操作
365
      const duplicatedMessages = await MessageModel.duplicateMessages(originalMessages);
366

367
      // 验证复制的消息数量是否正确
368
      expect(duplicatedMessages.length).toBe(originalMessages.length);
369

370
      // 验证每个复制的消息是否具有新的唯一ID，并且parentId被正确更新
371
      for (const original of originalMessages) {
372
        const copied = duplicatedMessages.find((m) => m.content === original.content);
373
        expect(copied).toBeDefined();
374
        expect(copied).not.toBeNull();
375
        expect(copied!.id).not.toBe(original.id);
376
        if (original.parentId) {
377
          const originalParent = originalMessages.find((m) => m.id === original.parentId);
378
          expect(originalParent).toBeDefined();
379
          const copiedParent = duplicatedMessages.find(
380
            (m) => m.content === originalParent!.content,
381
          );
382

383
          expect(copied!.parentId).toBe(copiedParent!.id);
384
        }
385
      }
386
    });
387
  });
388

389
  describe('clearTable', () => {
390
    it('should clear the table', async () => {
391
      await MessageModel.create(messageData);
392
      await MessageModel.clearTable();
393
      const messages = await MessageModel.queryAll();
394
      expect(messages).toHaveLength(0);
395
    });
396
  });
397

398
  describe('updatePluginState', () => {
399
    it('should update plugin state', async () => {
400
      const createdMessage = await MessageModel.create(messageData);
401
      await MessageModel.updatePluginState(createdMessage.id, { testKey: 'testValue' });
402
      const updatedMessage = await MessageModel.findById(createdMessage.id);
403
      expect(updatedMessage.pluginState).toHaveProperty('testKey', 'testValue');
404
    });
405
  });
406

407
  describe('updatePlugin', () => {
408
    it('should update plugin', async () => {
409
      const value = {
410
        identifier: 'testValue',
411
        arguments: 'abc',
412
        apiName: 'abc',
413
      };
414
      const createdMessage = await MessageModel.create(messageData);
415
      await MessageModel.updatePlugin(createdMessage.id, value);
416
      const updatedMessage = await MessageModel.findById(createdMessage.id);
417
      expect(updatedMessage.plugin).toEqual(value);
418
    });
419
  });
420

421
  describe('isEmpty', () => {
422
    it('should return true if table is empty', async () => {
423
      const number = await MessageModel.count();
424
      expect(number === 0).toBeTruthy();
425
    });
426
  });
427
});
428