lobe-chat

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

3
import { DBModel } from '@/database/client/core/types/db';
4
import { CreateMessageParams, MessageModel } from '@/database/client/models/message';
5
import { DB_Message } from '@/database/client/schemas/message';
6
import { DB_Topic } from '@/database/client/schemas/topic';
7
import { nanoid } from '@/utils/uuid';
8
import * as uuidUtils from '@/utils/uuid';
9

10
import { CreateTopicParams, QueryTopicParams, TopicModel } from '../topic';
11

12
describe('TopicModel', () => {
13
  let topicData: CreateTopicParams;
14
  const currentSessionId = 'session1';
15
  beforeEach(() => {
16
    // Set up topic data with the correct structure
17
    topicData = {
18
      sessionId: currentSessionId,
19
      title: 'Test Topic',
20
      favorite: false,
21
    };
22
  });
23

24
  afterEach(async () => {
25
    // Clean up the database after each test
26
    await TopicModel.clearTable();
27
  });
28

29
  describe('create', () => {
30
    it('should create a topic record', async () => {
31
      const result = await TopicModel.create(topicData);
32

33
      expect(result).toHaveProperty('id');
34
      // Verify that the topic has been added to the database
35
      const topicInDb = await TopicModel.findById(result.id);
36

37
      expect(topicInDb).toEqual(
38
        expect.objectContaining({
39
          title: topicData.title,
40
          favorite: topicData.favorite ? 1 : 0,
41
          sessionId: topicData.sessionId,
42
        }),
43
      );
44
    });
45

46
    it('should create a topic with favorite set to true', async () => {
47
      const favoriteTopicData: CreateTopicParams = {
48
        ...topicData,
49
        favorite: true,
50
      };
51
      const result = await TopicModel.create(favoriteTopicData);
52

53
      expect(result).toHaveProperty('id');
54
      const topicInDb = await TopicModel.findById(result.id);
55
      expect(topicInDb).toEqual(
56
        expect.objectContaining({
57
          title: favoriteTopicData.title,
58
          favorite: 1,
59
          sessionId: favoriteTopicData.sessionId,
60
        }),
61
      );
62
    });
63

64
    it('should update messages with the new topic id when messages are provided', async () => {
65
      const messagesToUpdate = [nanoid(), nanoid()];
66
      // 假设这些消息存在于数据库中
67
      for (const messageId of messagesToUpdate) {
68
        await MessageModel.table.add({ id: messageId, text: 'Sample message', topicId: null });
69
      }
70

71
      const topicDataWithMessages = {
72
        ...topicData,
73
        messages: messagesToUpdate,
74
      };
75

76
      const topic = await TopicModel.create(topicDataWithMessages);
77
      expect(topic).toHaveProperty('id');
78

79
      // 验证数据库中的消息是否已更新
80
      const updatedMessages: DB_Message[] = await MessageModel.table
81
        .where('id')
82
        .anyOf(messagesToUpdate)
83
        .toArray();
84

85
      expect(updatedMessages).toHaveLength(messagesToUpdate.length);
86
      for (const message of updatedMessages) {
87
        expect(message.topicId).toEqual(topic.id);
88
      }
89
    });
90

91
    it('should create a topic with a unique id when no id is provided', async () => {
92
      const spy = vi.spyOn(uuidUtils, 'nanoid'); // 使用 Vitest 的 spy 功能来监视 nanoid 调用
93
      const result = await TopicModel.create(topicData);
94

95
      expect(spy).toHaveBeenCalled(); // 验证 nanoid 被调用来生成 id
96
      expect(result).toHaveProperty('id');
97
      expect(typeof result.id).toBe('string');
98
      spy.mockRestore(); // 测试结束后恢复原始行为
99
    });
100
  });
101
  describe('batch create', () => {
102
    it('should batch create topic records', async () => {
103
      const topicsToCreate = [topicData, topicData];
104
      const results = await TopicModel.batchCreate(topicsToCreate);
105

106
      expect(results.ids).toHaveLength(topicsToCreate.length);
107
      // Verify that the topics have been added to the database
108
      for (const result of results.ids!) {
109
        const topicInDb = await TopicModel.findById(result);
110
        expect(topicInDb).toEqual(
111
          expect.objectContaining({
112
            title: topicData.title,
113
            favorite: topicData.favorite ? 1 : 0,
114
            sessionId: topicData.sessionId,
115
          }),
116
        );
117
      }
118
    });
119

120
    it('should batch create topics with mixed favorite values', async () => {
121
      const mixedTopicsData: CreateTopicParams[] = [
122
        { ...topicData, favorite: true },
123
        { ...topicData, favorite: false },
124
      ];
125

126
      const results = await TopicModel.batchCreate(mixedTopicsData);
127

128
      expect(results.ids).toHaveLength(mixedTopicsData.length);
129
      for (const id of results.ids!) {
130
        const topicInDb = await TopicModel.findById(id);
131
        expect(topicInDb).toBeDefined();
132
        expect(topicInDb.favorite).toBeGreaterThanOrEqual(0);
133
        expect(topicInDb.favorite).toBeLessThanOrEqual(1);
134
      }
135
    });
136
  });
137

138
  it('should query topics with pagination', async () => {
139
    // Create multiple topics to test the query method
140
    await TopicModel.batchCreate([topicData, topicData]);
141

142
    const queryParams: QueryTopicParams = { pageSize: 1, current: 0, sessionId: 'session1' };
143
    const queriedTopics = await TopicModel.query(queryParams);
144

145
    expect(queriedTopics).toHaveLength(1);
146
  });
147

148
  it('should find topics by session id', async () => {
149
    // Create multiple topics to test the findBySessionId method
150
    await TopicModel.batchCreate([topicData, topicData]);
151

152
    const topicsBySessionId = await TopicModel.findBySessionId(topicData.sessionId);
153

154
    expect(topicsBySessionId).toHaveLength(2);
155
    expect(topicsBySessionId.every((i) => i.sessionId === topicData.sessionId)).toBeTruthy();
156
  });
157

158
  it('should delete a topic and its associated messages', async () => {
159
    const createdTopic = await TopicModel.create(topicData);
160

161
    await TopicModel.delete(createdTopic.id);
162

163
    // Verify the topic and its related messages are deleted
164
    const topicInDb = await TopicModel.findById(createdTopic.id);
165
    expect(topicInDb).toBeUndefined();
166

167
    // You need to verify that messages related to the topic are also deleted
168
    // This would require additional setup to create messages associated with the topic
169
    // and then assertions to check that they're deleted after the topic itself is deleted
170
  });
171

172
  it('should batch delete topics by session id', async () => {
173
    // Create multiple topics to test the batchDeleteBySessionId method
174
    await TopicModel.batchCreate([topicData, topicData]);
175

176
    await TopicModel.batchDeleteBySessionId(topicData.sessionId);
177

178
    // Verify that all topics with the given session id are deleted
179
    const topicsInDb = await TopicModel.findBySessionId(topicData.sessionId);
180
    expect(topicsInDb).toHaveLength(0);
181
  });
182

183
  it('should update a topic', async () => {
184
    const createdTopic = await TopicModel.create(topicData);
185
    const updateData = { title: 'New Title' };
186

187
    await TopicModel.update(createdTopic.id, updateData);
188
    const updatedTopic = await TopicModel.findById(createdTopic.id);
189

190
    expect(updatedTopic).toHaveProperty('title', 'New Title');
191
  });
192

193
  describe('toggleFavorite', () => {
194
    it('should toggle favorite status of a topic', async () => {
195
      const createdTopic = await TopicModel.create(topicData);
196

197
      const newState = await TopicModel.toggleFavorite(createdTopic.id);
198

199
      expect(newState).toBe(true);
200
      const topicInDb = await TopicModel.findById(createdTopic.id);
201
      expect(topicInDb).toHaveProperty('favorite', 1);
202
    });
203

204
    it('should handle toggleFavorite when topic does not exist', async () => {
205
      const nonExistentTopicId = 'non-existent-id';
206
      await expect(TopicModel.toggleFavorite(nonExistentTopicId)).rejects.toThrow(
207
        `Topic with id ${nonExistentTopicId} not found`,
208
      );
209
    });
210

211
    it('should set favorite to specific state using toggleFavorite', async () => {
212
      const createdTopic = await TopicModel.create(topicData);
213

214
      // Set favorite to true regardless of current state
215
      await TopicModel.toggleFavorite(createdTopic.id, true);
216
      let topicInDb = await TopicModel.findById(createdTopic.id);
217
      expect(topicInDb.favorite).toBe(1);
218

219
      // Set favorite to false regardless of current state
220
      await TopicModel.toggleFavorite(createdTopic.id, false);
221
      topicInDb = await TopicModel.findById(createdTopic.id);
222
      expect(topicInDb.favorite).toBe(0);
223
    });
224
  });
225

226
  it('should delete a topic and its associated messages', async () => {
227
    // 创建话题和相关联的消息
228
    const createdTopic = await TopicModel.create(topicData);
229
    const messageData: CreateMessageParams = {
230
      content: 'Test Message',
231
      topicId: createdTopic.id,
232
      sessionId: topicData.sessionId,
233
      role: 'user',
234
    };
235
    await MessageModel.create(messageData);
236

237
    // 删除话题
238
    await TopicModel.delete(createdTopic.id);
239

240
    // 验证话题是否被删除
241
    const topicInDb = await TopicModel.findById(createdTopic.id);
242
    expect(topicInDb).toBeUndefined();
243

244
    // 验证与话题关联的消息是否也被删除
245
    const messagesInDb = await MessageModel.query({
246
      sessionId: topicData.sessionId,
247
      topicId: createdTopic.id,
248
    });
249
    expect(messagesInDb).toHaveLength(0);
250
  });
251

252
  it('should batch delete topics and their associated messages', async () => {
253
    // 创建多个话题和相关联的消息
254
    const createdTopic1 = await TopicModel.create(topicData);
255
    const createdTopic2 = await TopicModel.create(topicData);
256

257
    const messageData1: CreateMessageParams = {
258
      content: 'Test Message 1',
259
      topicId: createdTopic1.id,
260
      sessionId: topicData.sessionId,
261
      role: 'user',
262
    };
263
    const messageData2: CreateMessageParams = {
264
      content: 'Test Message 2',
265
      topicId: createdTopic2.id,
266
      sessionId: topicData.sessionId,
267
      role: 'user',
268
    };
269
    await MessageModel.create(messageData1);
270
    await MessageModel.create(messageData2);
271

272
    // 执行批量删除
273
    await TopicModel.batchDelete([createdTopic1.id, createdTopic2.id]);
274

275
    // 验证话题是否被删除
276
    const topicInDb1 = await TopicModel.findById(createdTopic1.id);
277
    const topicInDb2 = await TopicModel.findById(createdTopic2.id);
278
    expect(topicInDb1).toBeUndefined();
279
    expect(topicInDb2).toBeUndefined();
280

281
    // 验证与话题关联的消息是否也被删除
282
    const messagesInDb1 = await MessageModel.query({
283
      sessionId: topicData.sessionId,
284
      topicId: createdTopic1.id,
285
    });
286
    const messagesInDb2 = await MessageModel.query({
287
      sessionId: topicData.sessionId,
288
      topicId: createdTopic2.id,
289
    });
290
    expect(messagesInDb1).toHaveLength(0);
291
    expect(messagesInDb2).toHaveLength(0);
292
  });
293

294
  describe('duplicateTopic', () => {
295
    let originalTopic: DBModel<DB_Topic>;
296
    let originalMessages: any[];
297

298
    beforeEach(async () => {
299
      // 创建一个原始主题
300
      const { id } = await TopicModel.create({
301
        title: 'Original Topic',
302
        sessionId: 'session1',
303
        favorite: false,
304
      });
305
      originalTopic = await TopicModel.findById(id);
306

307
      // 创建一些关联到原始主题的消息
308
      originalMessages = await Promise.all(
309
        ['Message 1', 'Message 2'].map((text) =>
310
          MessageModel.create({
311
            content: text,
312
            topicId: originalTopic.id,
313
            sessionId: originalTopic.sessionId!,
314
            role: 'user',
315
          }),
316
        ),
317
      );
318
    });
319

320
    afterEach(async () => {
321
      // 清理数据库中的所有主题和消息
322
      await TopicModel.clearTable();
323
      await MessageModel.clearTable();
324
    });
325

326
    it('should duplicate a topic with all associated messages', async () => {
327
      // 执行复制操作
328
      await TopicModel.duplicateTopic(originalTopic.id);
329

330
      // 验证复制后的主题是否存在
331
      const duplicatedTopic = await TopicModel.findBySessionId(originalTopic.sessionId!);
332
      expect(duplicatedTopic).toHaveLength(2);
333

334
      // 验证复制后的消息是否存在
335
      const duplicatedMessages = await MessageModel.query({
336
        sessionId: originalTopic.sessionId!,
337
        topicId: duplicatedTopic[1].id, // 假设复制的主题是第二个
338
      });
339
      expect(duplicatedMessages).toHaveLength(originalMessages.length);
340
    });
341

342
    it('should throw an error if the topic does not exist', async () => {
343
      // 尝试复制一个不存在的主题
344
      const nonExistentTopicId = nanoid();
345
      await expect(TopicModel.duplicateTopic(nonExistentTopicId)).rejects.toThrow(
346
        `Topic with id ${nonExistentTopicId} not found`,
347
      );
348
    });
349

350
    it('should preserve the properties of the duplicated topic', async () => {
351
      // 执行复制操作
352
      await TopicModel.duplicateTopic(originalTopic.id);
353

354
      // 获取复制的主题
355
      const topics = await TopicModel.findBySessionId(originalTopic.sessionId!);
356
      const duplicatedTopic = topics.find((topic) => topic.id !== originalTopic.id);
357

358
      // 验证复制的主题是否保留了原始主题的属性
359
      expect(duplicatedTopic).toBeDefined();
360
      expect(duplicatedTopic).toMatchObject({
361
        title: originalTopic.title,
362
        favorite: originalTopic.favorite,
363
        sessionId: originalTopic.sessionId,
364
      });
365
      // 确保生成了新的 ID
366
      expect(duplicatedTopic.id).not.toBe(originalTopic.id);
367
    });
368

369
    it('should properly handle the messages hierarchy when duplicating', async () => {
370
      // 创建一个子消息关联到其中一个原始消息
371
      const { id } = await MessageModel.create({
372
        content: 'Child Message',
373
        topicId: originalTopic.id,
374
        parentId: originalMessages[0].id,
375
        sessionId: originalTopic.sessionId!,
376
        role: 'user',
377
      });
378
      const childMessage = await MessageModel.findById(id);
379

380
      // 执行复制操作
381
      await TopicModel.duplicateTopic(originalTopic.id);
382

383
      // 获取复制的消息
384
      const duplicatedMessages = await MessageModel.queryBySessionId(originalTopic.sessionId!);
385

386
      // 验证复制的子消息是否存在并且 parentId 已更新
387
      const duplicatedChildMessage = duplicatedMessages.find(
388
        (message) => message.content === childMessage.content && message.id !== childMessage.id,
389
      );
390

391
      expect(duplicatedChildMessage).toBeDefined();
392
      expect(duplicatedChildMessage.parentId).not.toBe(childMessage.parentId);
393
      expect(duplicatedChildMessage.parentId).toBeDefined();
394
    });
395

396
    it('should fail if the database transaction fails', async () => {
397
      // 强制数据库事务失败，例如通过在复制过程中抛出异常
398
      const dbTransactionFailedError = new Error('DB transaction failed');
399
      const spyOn = vi.spyOn(TopicModel['db'], 'transaction').mockImplementation((async () => {
400
        throw dbTransactionFailedError;
401
      }) as any);
402

403
      // 尝试复制主题并捕捉期望的错误
404
      await expect(TopicModel.duplicateTopic(originalTopic.id)).rejects.toThrow(
405
        dbTransactionFailedError,
406
      );
407
      spyOn.mockRestore();
408
    });
409

410
    it('should not create partial duplicates if the process fails at some point', async () => {
411
      // 假设复制消息的过程中发生了错误
412
      vi.spyOn(MessageModel, 'duplicateMessages').mockImplementation(async () => {
413
        throw new Error('Failed to duplicate messages');
414
      });
415

416
      // 尝试复制主题，期望会抛出错误
417
      await expect(TopicModel.duplicateTopic(originalTopic.id)).rejects.toThrow();
418

419
      // 确保没有创建任何副本
420
      const topics = await TopicModel.findBySessionId(originalTopic.sessionId!);
421
      expect(topics).toHaveLength(1); // 只有原始主题
422

423
      const messages = await MessageModel.queryBySessionId(originalTopic.sessionId!);
424
      expect(messages).toHaveLength(originalMessages.length); // 只有原始消息
425
    });
426
  });
427

428
  describe('clearTable', () => {
429
    it('should clear the table', async () => {
430
      // Create a topic to ensure the table is not empty
431
      await TopicModel.create(topicData);
432

433
      // Clear the table
434
      await TopicModel.clearTable();
435

436
      // Verify the table is empty
437
      const topics = await TopicModel.queryAll();
438
      expect(topics).toHaveLength(0);
439
    });
440
  });
441

442
  describe('update', () => {
443
    it('should update a topic', async () => {
444
      // Create a topic
445
      const createdTopic = await TopicModel.create(topicData);
446

447
      // Update the topic
448
      const newTitle = 'Updated Title';
449
      await TopicModel.update(createdTopic.id, { title: newTitle });
450

451
      // Verify the topic is updated
452
      const updatedTopic = await TopicModel.findById(createdTopic.id);
453
      expect(updatedTopic.title).toBe(newTitle);
454
    });
455
  });
456

457
  describe('batchDelete', () => {
458
    it('should batch delete topics', async () => {
459
      // Create multiple topics
460
      const topic1 = await TopicModel.create(topicData);
461
      const topic2 = await TopicModel.create(topicData);
462

463
      await TopicModel.create(topicData);
464

465
      const ids = [topic1.id, topic2.id];
466
      // Batch delete the topics
467
      await TopicModel.batchDelete(ids);
468

469
      expect(await TopicModel.table.count()).toEqual(1);
470
    });
471
  });
472

473
  describe('queryAll', () => {
474
    it('should query all topics', async () => {
475
      // Create multiple topics
476
      await TopicModel.batchCreate([topicData, topicData]);
477

478
      // Query all topics
479
      const topics = await TopicModel.queryAll();
480

481
      // Verify all topics are queried
482
      expect(topics).toHaveLength(2);
483
    });
484
  });
485

486
  describe('queryByKeyword', () => {
487
    it('should query global topics by keyword', async () => {
488
      // Create a topic with a unique title
489
      const uniqueTitle = 'Unique Title';
490
      await TopicModel.create({ ...topicData, title: uniqueTitle });
491

492
      // Query topics by the unique title
493
      const topics = await TopicModel.queryByKeyword(uniqueTitle);
494

495
      // Verify the correct topic is queried
496
      expect(topics).toHaveLength(1);
497
      expect(topics[0].title).toBe(uniqueTitle);
498
    });
499
    it('should query topics in current session by keyword', async () => {
500
      // Create a topic with a unique title
501
      const uniqueTitle = 'Unique Title';
502
      await TopicModel.create({ ...topicData, title: uniqueTitle });
503

504
      // Query topics by the unique title
505
      const topics = await TopicModel.queryByKeyword(uniqueTitle, currentSessionId);
506

507
      // Verify the correct topic is queried
508
      expect(topics).toHaveLength(1);
509
      expect(topics[0].title).toBe(uniqueTitle);
510
    });
511
    it('should not query any topic in other session by keyword', async () => {
512
      // Create a topic with a unique title
513
      const uniqueTitle = 'Unique Title';
514
      await TopicModel.create({ ...topicData, title: uniqueTitle });
515

516
      // Query topics by the unique title
517
      const topics = await TopicModel.queryByKeyword(uniqueTitle, 'session-id-2');
518

519
      // Verify the correct topic is queried
520
      expect(topics).toHaveLength(0);
521
    });
522
  });
523
});
524