UNPKG

@lobehub/chat

Version:

Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.

414 lines (339 loc) 12.4 kB
import { eq, not } from 'drizzle-orm/expressions'; import { Mock, beforeEach, describe, expect, it, vi } from 'vitest'; import { INBOX_SESSION_ID } from '@/const/session'; import { clientDB, initializeDB } from '@/database/client/db'; import { NewSession, SessionItem, agents, agentsToSessions, sessionGroups, sessions, users, } from '@/database/schemas'; import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent'; import { LobeAgentSession, LobeSessionType, SessionGroups } from '@/types/session'; import { ClientService } from './client'; const userId = 'message-db'; const sessionService = new ClientService(userId); const mockSessionId = 'mock-session-id'; const mockAgentId = 'agent-id'; // Mock data beforeEach(async () => { await initializeDB(); // 在每个测试用例之前,清空表 await clientDB.transaction(async (trx) => { await trx.insert(users).values([{ id: userId }, { id: '456' }]); await trx.insert(sessions).values([{ id: mockSessionId, userId }]); await trx.insert(agents).values([{ id: mockAgentId, userId }]); await trx .insert(agentsToSessions) .values([{ agentId: mockAgentId, sessionId: mockSessionId, userId }]); await trx.insert(sessionGroups).values([ { id: 'group-1', name: 'group-A', sort: 2, userId }, { id: 'group-2', name: 'group-B', sort: 1, userId }, { id: 'group-4', name: 'group-C', sort: 1, userId: '456' }, ]); }); }); afterEach(async () => { // 在每个测试用例之后,清空表 await clientDB.delete(users); }); describe('SessionService', () => { const mockSession = { id: mockSessionId, type: 'agent', meta: { title: 'Mock Session' }, } as LobeAgentSession; describe('createSession', () => { it('should create a new session and return its id', async () => { // Setup const sessionType = LobeSessionType.Agent; const defaultValue = { meta: { title: 'New Session' } } as Partial<LobeAgentSession>; // Execute const sessionId = await sessionService.createSession(sessionType, defaultValue); // Assert expect(sessionId).toMatch(/^ssn_/); }); }); describe('removeSession', () => { it('should remove a session by its id', async () => { // Execute await sessionService.removeSession(mockSessionId); // Assert const result = await clientDB.query.sessions.findFirst({ where: eq(sessions.id, mockSessionId), }); // Assert expect(result).toBeUndefined(); }); }); describe('removeAllSessions', () => { it('should clear all sessions from the table', async () => { // Setup await clientDB .insert(sessions) .values([{ userId: userId }, { userId: userId }, { userId: userId }]); // Execute await sessionService.removeAllSessions(); // Assert const result = await clientDB.query.sessions.findMany({ where: eq(sessionGroups.userId, userId), }); expect(result.length).toBe(0); }); }); describe('updateSession', () => { it('should update the group of a session', async () => { // Setup const groupId = 'group-1'; // Execute await sessionService.updateSession(mockSessionId, { group: groupId }); // Assert const result = await clientDB.query.sessions.findFirst({ where: eq(sessions.id, mockSessionId), }); expect(result).toMatchObject({ groupId }); }); it('should update the pinned status of a session', async () => { // Setup const pinned = true; // Execute await sessionService.updateSession(mockSessionId, { pinned }); // Assert const result = await clientDB.query.sessions.findFirst({ where: eq(sessions.id, mockSessionId), }); expect(result!.pinned).toBeTruthy(); }); }); describe('updateSessionConfig', () => { it('should update the config of a session', async () => { // Setup const newConfig = { model: 'abc' } as LobeAgentConfig; // Execute await sessionService.updateSessionConfig(mockSessionId, newConfig); // Assert const result = await sessionService.getSessionConfig(mockSessionId); expect(result).toMatchObject(newConfig); }); }); describe('countSessions', () => { it('should return false if no sessions exist', async () => { await clientDB.delete(sessions); // Execute const result = await sessionService.countSessions(); // Assert expect(result).toBe(0); }); it('should return true if sessions exist', async () => { // Setup await clientDB.delete(sessions); await clientDB.insert(sessions).values([{ userId }]); // Execute const result = await sessionService.countSessions(); // Assert expect(result).toBe(1); }); }); describe('searchSessions', () => { it('should return sessions that match the keyword', async () => { // Setup await clientDB.insert(agents).values({ userId, id: 'agent-1', title: 'Session Name' }); await clientDB .insert(agentsToSessions) .values({ agentId: 'agent-1', sessionId: mockSessionId, userId }); // Execute const keyword = 'Name'; const result = await sessionService.searchSessions(keyword); // Assert // TODO: 后续需要把这个搜索的标题和描述都加上,现在这个 client 搜索会有问题 expect(result).toMatchObject([{ id: mockSessionId }]); }); }); describe('cloneSession', () => { it('should duplicate a session and return its id', async () => { // Setup const newTitle = 'Duplicated Session'; const session: NewSession = { id: 'duplicated-session-id', title: '123', userId, }; await clientDB.insert(sessions).values([session]); await clientDB.insert(agents).values({ userId, id: 'agent-1' }); await clientDB .insert(agentsToSessions) .values({ agentId: 'agent-1', sessionId: 'duplicated-session-id', userId }); // Execute const duplicatedSessionId = await sessionService.cloneSession(mockSessionId, newTitle); // Assert const result = await clientDB.query.sessions.findFirst({ where: eq(sessions.id, duplicatedSessionId!), }); expect(result).toMatchObject({ title: 'Duplicated Session' }); }); }); describe('getGroupedSessions', () => { it('should retrieve sessions with their group', async () => { // Execute const sessionsWithGroup = await sessionService.getGroupedSessions(); expect(sessionsWithGroup).toMatchObject({ sessionGroups: [ { id: 'group-2', name: 'group-B', sort: 1 }, { id: 'group-1', name: 'group-A', sort: 2 }, ], sessions: [{ id: 'mock-session-id', type: 'agent' }], }); }); }); describe('getSessionsByType', () => { it('should get sessions by type "all"', async () => { const sessions = await sessionService.getSessionsByType('all'); expect(sessions).toBeDefined(); }); it('should get sessions by type "agent"', async () => { const sessions = await sessionService.getSessionsByType('agent'); expect(sessions).toBeDefined(); }); it('should get sessions by type "group"', async () => { const sessions = await sessionService.getSessionsByType('group'); expect(sessions).toBeDefined(); }); }); describe('getSessionConfig', () => { it.skip('should get default config for INBOX_SESSION_ID', async () => { const config = await sessionService.getSessionConfig(INBOX_SESSION_ID); expect(config).toBeDefined(); }); it('should throw error for non-existent session', async () => { await expect(sessionService.getSessionConfig('non-existent')).rejects.toThrow( 'Session not found', ); }); }); describe('updateSessionMeta', () => { it('should not update meta for INBOX_SESSION_ID', async () => { const result = await sessionService.updateSessionMeta(INBOX_SESSION_ID, { title: 'New Title', }); expect(result).toBeUndefined(); }); it('should update meta for normal session', async () => { const meta = { title: 'Updated Title' }; await sessionService.updateSessionMeta(mockSessionId, meta); const session = await clientDB.query.sessions.findFirst({ where: eq(sessions.id, mockSessionId), }); expect(session).toBeDefined(); }); }); describe('updateSessionChatConfig', () => { it('should update chat config', async () => { const chatConfig = { temperature: 0.8 } as Partial<LobeAgentChatConfig>; const result = await sessionService.updateSessionChatConfig(mockSessionId, chatConfig); expect(result).toBeDefined(); }); }); describe('model getters', () => { it('should return session model instance', () => { // @ts-ignore - accessing private getter const model = sessionService.sessionModel; expect(model).toBeDefined(); }); it('should return session group model instance', () => { // @ts-ignore - accessing private getter const model = sessionService.sessionGroupModel; expect(model).toBeDefined(); }); }); // SessionGroup related tests describe('createSessionGroup', () => { it('should create a new session group and return its id', async () => { // Setup const groupName = 'New Group'; const sort = 1; // Execute const groupId = await sessionService.createSessionGroup(groupName, sort); // Assert expect(groupId).toMatch(/^sg_/); const result = await clientDB.query.sessionGroups.findFirst({ where: eq(sessionGroups.id, groupId), }); expect(result).toMatchObject({ id: groupId, name: groupName, sort }); }); }); describe('removeSessionGroup', () => { it('should remove a session group by its id', async () => { const groupId = 'group-1'; // Execute await sessionService.removeSessionGroup(groupId); const result = await clientDB.query.sessionGroups.findFirst({ where: eq(sessionGroups.id, groupId), }); // Assert expect(result).toBeUndefined(); }); }); describe('clearSessionGroups', () => { it('should clear all session groups', async () => { // Execute await sessionService.removeSessionGroups(); // Assert const result = await clientDB.query.sessionGroups.findMany({ where: eq(sessionGroups.userId, userId), }); expect(result.length).toBe(0); const result2 = await clientDB.query.sessionGroups.findMany({ where: not(eq(sessionGroups.userId, userId)), }); expect(result2.length).toBeGreaterThan(0); }); }); describe('getSessionGroups', () => { it('should retrieve all session groups', async () => { // Execute const result = await sessionService.getSessionGroups(); // Assert const groups = [ { id: 'group-2', name: 'group-B', sort: 1 }, { id: 'group-1', name: 'group-A', sort: 2 }, ]; expect(result).toMatchObject(groups); }); }); describe('updateSessionGroup', () => { it('should update a session group', async () => { // Setup const groupId = 'group-1'; const data = { name: 'Updated Group', sort: 2 }; // Execute await sessionService.updateSessionGroup(groupId, data); // Assert const result = await clientDB.query.sessionGroups.findFirst({ where: eq(sessionGroups.id, groupId), }); expect(result).toMatchObject({ id: groupId, ...data }); }); }); describe('updateSessionGroupOrder', () => { it('should update the order of session groups', async () => { // Setup const sortMap = [ { id: 'group-1', sort: 2 }, { id: 'group-2', sort: 1 }, ]; // Execute await sessionService.updateSessionGroupOrder(sortMap); // Assert const data = await clientDB.query.sessionGroups.findMany({ where: eq(sessionGroups.userId, userId), }); expect(data).toMatchObject([ { id: 'group-1', sort: 2 }, { id: 'group-2', sort: 1 }, ]); }); }); });