lobe-chat

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

3
import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
4
import { CreateMessageParams, MessageModel } from '@/database/client/models/message';
5
import { SessionGroupModel } from '@/database/client/models/sessionGroup';
6
import { TopicModel } from '@/database/client/models/topic';
7
import { LobeAgentConfig } from '@/types/agent';
8
import {
9
  LobeAgentSession,
10
  LobeSessionType,
11
  SessionDefaultGroup,
12
  SessionGroupId,
13
} from '@/types/session';
14

15
import { SessionModel } from '../session';
16

17
describe('SessionModel', () => {
18
  let sessionData: Partial<LobeAgentSession>;
19

20
  beforeEach(() => {
21
    // Set up session data with the correct structure
22
    sessionData = {
23
      type: LobeSessionType.Agent,
24
      group: 'testGroup',
25
      meta: {},
26
      config: DEFAULT_AGENT_CONFIG,
27
      // ... other properties based on LobeAgentSession
28
    };
29
  });
30

31
  afterEach(async () => {
32
    // Clean up the database after each test
33
    await SessionModel.clearTable();
34
  });
35
  describe('create', () => {
36
    it('should create a session record', async () => {
37
      const result = await SessionModel.create('agent', sessionData);
38

39
      expect(result).toHaveProperty('id');
40
      // Verify that the session has been added to the database
41
      // Assuming findById is a method that retrieves a session by ID
42
      const sessionInDb = await SessionModel.findById(result.id);
43

44
      expect(sessionInDb).toEqual(expect.objectContaining(sessionData));
45
    });
46
  });
47

48
  describe('batchCreate', () => {
49
    it('should batch create session records', async () => {
50
      const sessionsToCreate = [sessionData, sessionData];
51
      const results = await SessionModel.batchCreate(sessionsToCreate as LobeAgentSession[]);
52

53
      expect(results.ids).toHaveLength(sessionsToCreate.length);
54
      // Verify that the sessions have been added to the database
55
      for (const result of results.ids!) {
56
        const sessionInDb = await SessionModel.findById(result);
57
        expect(sessionInDb).toEqual(expect.objectContaining(sessionData));
58
      }
59
    });
60

61
    it('should set group to default if it does not exist in SessionGroup', async () => {
62
      const sessionDataWithInvalidGroup = {
63
        ...sessionData,
64
        group: 'nonExistentGroup',
65
      } as LobeAgentSession;
66

67
      const results = await SessionModel.batchCreate([sessionDataWithInvalidGroup]);
68

69
      // Verify that the group has been set to default
70
      for (const result of results.ids!) {
71
        const sessionInDb = await SessionModel.findById(result);
72
        expect(sessionInDb.group).toEqual(SessionDefaultGroup.Default);
73
      }
74
    });
75
  });
76

77
  describe('query', () => {
78
    it('should query sessions with pagination', async () => {
79
      // Create multiple sessions to test the query method
80
      await SessionModel.batchCreate([sessionData, sessionData] as LobeAgentSession[]);
81

82
      const queriedSessions = await SessionModel.query({ pageSize: 1, current: 0 });
83

84
      expect(queriedSessions).toHaveLength(1);
85
    });
86
  });
87

88
  describe('querySessionsByGroupId', () => {
89
    it('should query sessions by group', async () => {
90
      // Create multiple sessions to test the queryByGroup method
91
      const group: SessionGroupId = 'testGroup';
92
      await SessionGroupModel.create('测试分组', 0, group);
93

94
      await SessionModel.batchCreate([sessionData, sessionData] as LobeAgentSession[]);
95

96
      const sessionsByGroup = await SessionModel.querySessionsByGroupId(group);
97

98
      // Assuming all created sessions belong to the same group
99
      expect(sessionsByGroup).toHaveLength(2);
100
      expect(sessionsByGroup.every((i) => i.group === group)).toBeTruthy();
101
    });
102
  });
103

104
  describe('update', () => {
105
    it('should update a session', async () => {
106
      const createdSession = await SessionModel.create('agent', sessionData);
107
      const updateData = { group: 'newGroup' };
108

109
      await SessionModel.update(createdSession.id, updateData);
110
      const updatedSession = await SessionModel.findById(createdSession.id);
111

112
      expect(updatedSession).toHaveProperty('group', 'newGroup');
113
    });
114

115
    it('should update pinned status of a session', async () => {
116
      const createdSession = await SessionModel.create('agent', sessionData);
117
      await SessionModel.update(createdSession.id, { pinned: 1 });
118
      const updatedSession = await SessionModel.findById(createdSession.id);
119
      expect(updatedSession).toHaveProperty('pinned', 1);
120
    });
121
  });
122

123
  describe('updateConfig', () => {
124
    it('should update config of a session', async () => {
125
      const { id } = await SessionModel.create('agent', sessionData);
126
      const dbSession = await SessionModel.findById(id);
127

128
      const newConfig = { ...dbSession.config, systemRole: 'newValue' } as LobeAgentConfig;
129
      await SessionModel.updateConfig(id, newConfig);
130
      const updatedSession = await SessionModel.findById(id);
131
      expect(updatedSession.config).toEqual(newConfig);
132
    });
133
  });
134

135
  describe('clearTable', () => {
136
    it('should clear all sessions', async () => {
137
      await SessionModel.batchCreate([sessionData, sessionData] as LobeAgentSession[]);
138
      await SessionModel.clearTable();
139
      const sessionsInDb = await SessionModel.query();
140
      expect(sessionsInDb).toHaveLength(0);
141
    });
142
  });
143

144
  describe('isEmpty', () => {
145
    it('should check if table is empty', async () => {
146
      await SessionModel.clearTable();
147
      const isEmpty = await SessionModel.isEmpty();
148
      expect(isEmpty).toBeTruthy();
149
      await SessionModel.create('agent', sessionData);
150
      const isNotEmpty = await SessionModel.isEmpty();
151
      expect(isNotEmpty).toBeFalsy();
152
    });
153
  });
154

155
  describe('queryByKeyword', () => {
156
    it('should query sessions by keyword', async () => {
157
      const keyword = 'testKeyword';
158
      const sessionWithKeyword = { ...sessionData, meta: { title: keyword } };
159
      await SessionModel.create('agent', sessionWithKeyword);
160
      const sessionsByKeyword = await SessionModel.queryByKeyword(keyword);
161
      expect(sessionsByKeyword).toHaveLength(1);
162
      expect(sessionsByKeyword[0].meta.title).toContain(keyword);
163
    });
164
  });
165

166
  describe('getPinnedSessions', () => {
167
    it('should get pinned sessions', async () => {
168
      const pinnedSession = { ...sessionData, pinned: true };
169
      const unpinnedSession = { ...sessionData, pinned: false };
170
      await SessionModel.batchCreate([pinnedSession, unpinnedSession] as LobeAgentSession[]);
171
      const pinnedSessions = await SessionModel.getPinnedSessions();
172
      expect(pinnedSessions).toHaveLength(1);
173
      expect(pinnedSessions[0].pinned).toBeTruthy();
174
    });
175
  });
176

177
  describe('queryWithGroups', () => {
178
    it('should query sessions with groups', async () => {
179
      await SessionModel.create('agent', sessionData);
180

181
      const sessionsWithGroups = await SessionModel.queryWithGroups();
182
      expect(sessionsWithGroups.sessions).toHaveLength(1);
183
      expect(sessionsWithGroups.sessions[0]).toEqual(expect.objectContaining(sessionData));
184
    });
185
  });
186

187
  describe('queryByGroupIds', () => {
188
    it('should query sessions by group ids', async () => {
189
      const createdSession = await SessionModel.create('agent', sessionData);
190
      const session = await SessionModel.findById(createdSession.id);
191
      const sessionsByGroupIds = await SessionModel.queryByGroupIds([session.group]);
192
      expect(sessionsByGroupIds[session.group]).toHaveLength(1);
193
      expect(sessionsByGroupIds[session.group][0]).toEqual(expect.objectContaining(sessionData));
194
    });
195
  });
196

197
  describe('delete', () => {
198
    it('should delete a session', async () => {
199
      const createdSession = await SessionModel.create('agent', sessionData);
200
      await SessionModel.delete(createdSession.id);
201
      const sessionInDb = await SessionModel.findById(createdSession.id);
202
      expect(sessionInDb).toBeUndefined();
203
    });
204

205
    // 删除一个 session 时，也需要同步删除具有 sessionId 的 topic 和 message
206
    it('should delete a session and its associated data', async () => {
207
      // create session , topic and message test data
208
      const { id: sessionId } = await SessionModel.create('agent', sessionData);
209

210
      const topicData = {
211
        title: 'Test Topic',
212
        sessionId: sessionId,
213
        favorite: false,
214
      };
215
      const createdTopic = await TopicModel.create(topicData);
216

217
      const messageData: CreateMessageParams = {
218
        content: 'Test Message',
219
        sessionId: sessionId,
220
        topicId: createdTopic.id,
221
        role: 'user',
222
      };
223
      await MessageModel.create(messageData);
224

225
      await SessionModel.delete(sessionId);
226

227
      // Verify the session and its related data (topics, messages) are deleted
228
      const sessionInDb = await SessionModel.findById(sessionId);
229
      expect(sessionInDb).toBeUndefined();
230

231
      // You need to verify that topics and messages related to the session are also deleted
232
      const topicsInDb = await TopicModel.findBySessionId(sessionId);
233
      expect(topicsInDb).toHaveLength(0);
234

235
      // Verify all associated messages are deleted
236
      const messagesInDb = await MessageModel.query({ sessionId });
237
      expect(messagesInDb).toHaveLength(0);
238
    });
239
  });
240

241
  describe('duplicate', () => {
242
    it('should duplicate a session', async () => {
243
      const createdSession = await SessionModel.create('agent', sessionData);
244
      const duplicatedSession = await SessionModel.duplicate(createdSession.id);
245
      const sessionsInDb = await SessionModel.query();
246
      expect(sessionsInDb).toHaveLength(2);
247

248
      if (!duplicatedSession) return;
249
      const session = await SessionModel.findById(duplicatedSession.id);
250

251
      expect(session).toEqual(expect.objectContaining(sessionData));
252
    });
253
  });
254
});
255