@lobehub/lobehub 2.0.0-next.137 → 2.0.0-next.139

package/packages/database/src/models/__tests__/thread.test.ts

@@ -0,0 +1,327 @@
+import { ThreadStatus, ThreadType } from '@lobechat/types';
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { sessions, threads, topics, users } from '../../schemas';
+import { LobeChatDatabase } from '../../type';
+import { ThreadModel } from '../thread';
+import { getTestDB } from './_util';
+
+const userId = 'thread-user-test';
+const otherUserId = 'other-user-test';
+const sessionId = 'thread-session';
+const topicId = 'thread-topic';
+
+const serverDB: LobeChatDatabase = await getTestDB();
+const threadModel = new ThreadModel(serverDB, userId);
+
+describe('ThreadModel', () => {
+  beforeEach(async () => {
+    await serverDB.delete(users);
+
+    // Create test users, session and topic
+    await serverDB.transaction(async (tx) => {
+      await tx.insert(users).values([{ id: userId }, { id: otherUserId }]);
+      await tx.insert(sessions).values({ id: sessionId, userId });
+      await tx.insert(topics).values({ id: topicId, userId, sessionId });
+    });
+  });
+
+  afterEach(async () => {
+    await serverDB.delete(users);
+  });
+
+  describe('create', () => {
+    it('should create a new thread', async () => {
+      const result = await threadModel.create({
+        topicId,
+        type: ThreadType.Standalone,
+        sourceMessageId: 'msg-1',
+      });
+
+      expect(result).toBeDefined();
+      expect(result.topicId).toBe(topicId);
+      expect(result.type).toBe(ThreadType.Standalone);
+      expect(result.status).toBe(ThreadStatus.Active);
+      expect(result.sourceMessageId).toBe('msg-1');
+    });
+
+    it('should create a thread with title', async () => {
+      const result = await threadModel.create({
+        topicId,
+        type: ThreadType.Continuation,
+        title: 'Test Thread',
+      });
+
+      expect(result.title).toBe('Test Thread');
+      expect(result.type).toBe(ThreadType.Continuation);
+    });
+  });
+
+  describe('query', () => {
+    it('should return all threads for the user', async () => {
+      // Create test threads
+      await serverDB.insert(threads).values([
+        {
+          id: 'thread-1',
+          topicId,
+          type: ThreadType.Standalone,
+          status: ThreadStatus.Active,
+          userId,
+          updatedAt: new Date('2024-01-01'),
+        },
+        {
+          id: 'thread-2',
+          topicId,
+          type: ThreadType.Continuation,
+          status: ThreadStatus.Active,
+          userId,
+          updatedAt: new Date('2024-01-02'),
+        },
+      ]);
+
+      const result = await threadModel.query();
+
+      expect(result).toHaveLength(2);
+      // Should be ordered by updatedAt desc
+      expect(result[0].id).toBe('thread-2');
+      expect(result[1].id).toBe('thread-1');
+    });
+
+    it('should only return threads for the current user', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'other-topic', userId: otherUserId });
+        await tx.insert(threads).values([
+          {
+            id: 'thread-1',
+            topicId,
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId,
+          },
+          {
+            id: 'thread-2',
+            topicId: 'other-topic',
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId: otherUserId,
+          },
+        ]);
+      });
+
+      const result = await threadModel.query();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('thread-1');
+    });
+  });
+
+  describe('queryByTopicId', () => {
+    it('should return threads for a specific topic', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'another-topic', userId, sessionId });
+        await tx.insert(threads).values([
+          {
+            id: 'thread-1',
+            topicId,
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId,
+            updatedAt: new Date('2024-01-01'),
+          },
+          {
+            id: 'thread-2',
+            topicId: 'another-topic',
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId,
+            updatedAt: new Date('2024-01-02'),
+          },
+        ]);
+      });
+
+      const result = await threadModel.queryByTopicId(topicId);
+
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('thread-1');
+    });
+
+    it('should return empty array when no threads exist for the topic', async () => {
+      const result = await threadModel.queryByTopicId('non-existent-topic');
+
+      expect(result).toHaveLength(0);
+    });
+  });
+
+  describe('findById', () => {
+    it('should return a thread by id', async () => {
+      await serverDB.insert(threads).values({
+        id: 'thread-1',
+        topicId,
+        type: ThreadType.Standalone,
+        status: ThreadStatus.Active,
+        userId,
+        title: 'Test Thread',
+      });
+
+      const result = await threadModel.findById('thread-1');
+
+      expect(result).toBeDefined();
+      expect(result?.id).toBe('thread-1');
+      expect(result?.title).toBe('Test Thread');
+    });
+
+    it('should return undefined for non-existent thread', async () => {
+      const result = await threadModel.findById('non-existent');
+
+      expect(result).toBeUndefined();
+    });
+
+    it('should not return thread belonging to another user', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'other-topic', userId: otherUserId });
+        await tx.insert(threads).values({
+          id: 'thread-other',
+          topicId: 'other-topic',
+          type: ThreadType.Standalone,
+          status: ThreadStatus.Active,
+          userId: otherUserId,
+        });
+      });
+
+      const result = await threadModel.findById('thread-other');
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('update', () => {
+    it('should update a thread', async () => {
+      await serverDB.insert(threads).values({
+        id: 'thread-1',
+        topicId,
+        type: ThreadType.Standalone,
+        status: ThreadStatus.Active,
+        userId,
+        title: 'Original Title',
+      });
+
+      await threadModel.update('thread-1', {
+        title: 'Updated Title',
+        status: ThreadStatus.Archived,
+      });
+
+      const updated = await serverDB.query.threads.findFirst({
+        where: eq(threads.id, 'thread-1'),
+      });
+
+      expect(updated?.title).toBe('Updated Title');
+      expect(updated?.status).toBe(ThreadStatus.Archived);
+    });
+
+    it('should not update thread belonging to another user', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'other-topic', userId: otherUserId });
+        await tx.insert(threads).values({
+          id: 'thread-other',
+          topicId: 'other-topic',
+          type: ThreadType.Standalone,
+          status: ThreadStatus.Active,
+          userId: otherUserId,
+          title: 'Original Title',
+        });
+      });
+
+      await threadModel.update('thread-other', { title: 'Hacked Title' });
+
+      const unchanged = await serverDB.query.threads.findFirst({
+        where: eq(threads.id, 'thread-other'),
+      });
+
+      expect(unchanged?.title).toBe('Original Title');
+    });
+  });
+
+  describe('delete', () => {
+    it('should delete a thread', async () => {
+      await serverDB.insert(threads).values({
+        id: 'thread-1',
+        topicId,
+        type: ThreadType.Standalone,
+        status: ThreadStatus.Active,
+        userId,
+      });
+
+      await threadModel.delete('thread-1');
+
+      const deleted = await serverDB.query.threads.findFirst({
+        where: eq(threads.id, 'thread-1'),
+      });
+
+      expect(deleted).toBeUndefined();
+    });
+
+    it('should not delete thread belonging to another user', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'other-topic', userId: otherUserId });
+        await tx.insert(threads).values({
+          id: 'thread-other',
+          topicId: 'other-topic',
+          type: ThreadType.Standalone,
+          status: ThreadStatus.Active,
+          userId: otherUserId,
+        });
+      });
+
+      await threadModel.delete('thread-other');
+
+      const stillExists = await serverDB.query.threads.findFirst({
+        where: eq(threads.id, 'thread-other'),
+      });
+
+      expect(stillExists).toBeDefined();
+    });
+  });
+
+  describe('deleteAll', () => {
+    it('should delete all threads for the current user', async () => {
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values({ id: 'other-topic', userId: otherUserId });
+        await tx.insert(threads).values([
+          {
+            id: 'thread-1',
+            topicId,
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId,
+          },
+          {
+            id: 'thread-2',
+            topicId,
+            type: ThreadType.Continuation,
+            status: ThreadStatus.Active,
+            userId,
+          },
+          {
+            id: 'thread-3',
+            topicId: 'other-topic',
+            type: ThreadType.Standalone,
+            status: ThreadStatus.Active,
+            userId: otherUserId,
+          },
+        ]);
+      });
+
+      await threadModel.deleteAll();
+
+      const userThreads = await serverDB.select().from(threads).where(eq(threads.userId, userId));
+      const otherUserThreads = await serverDB
+        .select()
+        .from(threads)
+        .where(eq(threads.userId, otherUserId));
+
+      expect(userThreads).toHaveLength(0);
+      expect(otherUserThreads).toHaveLength(1);
+    });
+  });
+});

package/packages/database/src/models/__tests__/user.test.ts

@@ -0,0 +1,372 @@
+import { UserPreference } from '@lobechat/types';
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { nextauthAccounts, userSettings, users } from '../../schemas';
+import { LobeChatDatabase } from '../../type';
+import { UserModel, UserNotFoundError } from '../user';
+import { getTestDB } from './_util';
+
+const userId = 'user-model-test';
+const otherUserId = 'other-user-test';
+
+const serverDB: LobeChatDatabase = await getTestDB();
+const userModel = new UserModel(serverDB, userId);
+
+// Mock decryptor function
+const mockDecryptor = vi.fn().mockResolvedValue({});
+
+describe('UserModel', () => {
+  beforeEach(async () => {
+    await serverDB.delete(users);
+    await serverDB.insert(users).values([
+      { id: userId, email: 'test@example.com', fullName: 'Test User' },
+      { id: otherUserId, email: 'other@example.com' },
+    ]);
+  });
+
+  afterEach(async () => {
+    await serverDB.delete(users);
+    vi.clearAllMocks();
+  });
+
+  describe('getUserRegistrationDuration', () => {
+    it('should return registration duration for existing user', async () => {
+      const thirtyDaysAgo = new Date();
+      thirtyDaysAgo.setDate(thirtyDaysAgo.getDate() - 30);
+
+      await serverDB.update(users).set({ createdAt: thirtyDaysAgo }).where(eq(users.id, userId));
+
+      const result = await userModel.getUserRegistrationDuration();
+
+      expect(result.duration).toBeGreaterThanOrEqual(30);
+      expect(result.createdAt).toBeDefined();
+      expect(result.updatedAt).toBeDefined();
+    });
+
+    it('should return default duration for non-existent user', async () => {
+      const nonExistentUserModel = new UserModel(serverDB, 'non-existent');
+
+      const result = await nonExistentUserModel.getUserRegistrationDuration();
+
+      expect(result.duration).toBe(1);
+    });
+  });
+
+  describe('getUserState', () => {
+    it('should return user state with settings', async () => {
+      // Create user settings
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        general: { fontSize: 14 },
+        tts: { voice: 'default' },
+      });
+
+      const result = await userModel.getUserState(mockDecryptor);
+
+      expect(result.userId).toBe(userId);
+      expect(result.email).toBe('test@example.com');
+      expect(result.fullName).toBe('Test User');
+      expect(result.settings.general).toEqual({ fontSize: 14 });
+      expect(result.settings.tts).toEqual({ voice: 'default' });
+    });
+
+    it('should throw UserNotFoundError for non-existent user', async () => {
+      const nonExistentUserModel = new UserModel(serverDB, 'non-existent');
+
+      await expect(nonExistentUserModel.getUserState(mockDecryptor)).rejects.toThrow(
+        UserNotFoundError,
+      );
+    });
+
+    it('should handle decryptor errors gracefully', async () => {
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        keyVaults: 'encrypted-data',
+      });
+
+      const failingDecryptor = vi.fn().mockRejectedValue(new Error('Decryption failed'));
+
+      const result = await userModel.getUserState(failingDecryptor);
+
+      expect(result.settings.keyVaults).toEqual({});
+    });
+  });
+
+  describe('getUserSSOProviders', () => {
+    it('should return SSO providers for user', async () => {
+      await serverDB.insert(nextauthAccounts).values({
+        userId,
+        provider: 'google',
+        providerAccountId: 'google-123',
+        type: 'oauth' as any,
+      });
+
+      const result = await userModel.getUserSSOProviders();
+
+      expect(result).toHaveLength(1);
+      expect(result[0].provider).toBe('google');
+      expect(result[0].providerAccountId).toBe('google-123');
+    });
+
+    it('should return empty array when no SSO providers', async () => {
+      const result = await userModel.getUserSSOProviders();
+
+      expect(result).toHaveLength(0);
+    });
+  });
+
+  describe('getUserSettings', () => {
+    it('should return user settings', async () => {
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        general: { fontSize: 14 },
+      });
+
+      const result = await userModel.getUserSettings();
+
+      expect(result).toBeDefined();
+      expect(result?.general).toEqual({ fontSize: 14 });
+    });
+
+    it('should return undefined when no settings exist', async () => {
+      const result = await userModel.getUserSettings();
+
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('updateUser', () => {
+    it('should update user properties', async () => {
+      await userModel.updateUser({
+        fullName: 'Updated Name',
+        avatar: 'https://example.com/avatar.jpg',
+      });
+
+      const updated = await serverDB.query.users.findFirst({
+        where: eq(users.id, userId),
+      });
+
+      expect(updated?.fullName).toBe('Updated Name');
+      expect(updated?.avatar).toBe('https://example.com/avatar.jpg');
+    });
+  });
+
+  describe('deleteSetting', () => {
+    it('should delete user settings', async () => {
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        general: { fontSize: 14 },
+      });
+
+      await userModel.deleteSetting();
+
+      const settings = await serverDB.query.userSettings.findFirst({
+        where: eq(userSettings.id, userId),
+      });
+
+      expect(settings).toBeUndefined();
+    });
+  });
+
+  describe('updateSetting', () => {
+    it('should create settings if not exist', async () => {
+      await userModel.updateSetting({
+        general: { fontSize: 16 },
+      });
+
+      const settings = await serverDB.query.userSettings.findFirst({
+        where: eq(userSettings.id, userId),
+      });
+
+      expect(settings?.general).toEqual({ fontSize: 16 });
+    });
+
+    it('should update existing settings', async () => {
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        general: { fontSize: 14 },
+      });
+
+      await userModel.updateSetting({
+        general: { fontSize: 18 },
+      });
+
+      const settings = await serverDB.query.userSettings.findFirst({
+        where: eq(userSettings.id, userId),
+      });
+
+      expect(settings?.general).toEqual({ fontSize: 18 });
+    });
+  });
+
+  describe('updatePreference', () => {
+    it('should update user preference', async () => {
+      await userModel.updatePreference({
+        telemetry: false,
+      });
+
+      const user = await serverDB.query.users.findFirst({
+        where: eq(users.id, userId),
+      });
+
+      expect((user?.preference as UserPreference)?.telemetry).toBe(false);
+    });
+
+    it('should merge with existing preference', async () => {
+      await serverDB
+        .update(users)
+        .set({ preference: { telemetry: true, useCmdEnterToSend: true } })
+        .where(eq(users.id, userId));
+
+      await userModel.updatePreference({
+        telemetry: false,
+      });
+
+      const user = await serverDB.query.users.findFirst({
+        where: eq(users.id, userId),
+      });
+
+      const preference = user?.preference as UserPreference;
+      expect(preference?.telemetry).toBe(false);
+      expect(preference?.useCmdEnterToSend).toBe(true);
+    });
+
+    it('should do nothing for non-existent user', async () => {
+      const nonExistentUserModel = new UserModel(serverDB, 'non-existent');
+
+      await expect(
+        nonExistentUserModel.updatePreference({ telemetry: false }),
+      ).resolves.toBeUndefined();
+    });
+  });
+
+  describe('updateGuide', () => {
+    it('should update user guide preference', async () => {
+      await userModel.updateGuide({
+        moveSettingsToAvatar: true,
+      });
+
+      const user = await serverDB.query.users.findFirst({
+        where: eq(users.id, userId),
+      });
+
+      const preference = user?.preference as UserPreference;
+      expect(preference?.guide?.moveSettingsToAvatar).toBe(true);
+    });
+
+    it('should do nothing for non-existent user', async () => {
+      const nonExistentUserModel = new UserModel(serverDB, 'non-existent');
+
+      await expect(
+        nonExistentUserModel.updateGuide({ moveSettingsToAvatar: true }),
+      ).resolves.toBeUndefined();
+    });
+  });
+
+  describe('static methods', () => {
+    describe('makeSureUserExist', () => {
+      it('should create user if not exists', async () => {
+        await UserModel.makeSureUserExist(serverDB, 'new-user-id');
+
+        const user = await serverDB.query.users.findFirst({
+          where: eq(users.id, 'new-user-id'),
+        });
+
+        expect(user).toBeDefined();
+      });
+
+      it('should not throw if user already exists', async () => {
+        await expect(UserModel.makeSureUserExist(serverDB, userId)).resolves.not.toThrow();
+      });
+    });
+
+    describe('createUser', () => {
+      it('should create a new user', async () => {
+        const result = await UserModel.createUser(serverDB, {
+          id: 'brand-new-user',
+          email: 'new@example.com',
+        });
+
+        expect(result.duplicate).toBe(false);
+        expect(result.user?.id).toBe('brand-new-user');
+        expect(result.user?.email).toBe('new@example.com');
+      });
+
+      it('should return duplicate flag for existing user', async () => {
+        const result = await UserModel.createUser(serverDB, {
+          id: userId,
+          email: 'duplicate@example.com',
+        });
+
+        expect(result.duplicate).toBe(true);
+      });
+    });
+
+    describe('deleteUser', () => {
+      it('should delete a user', async () => {
+        await UserModel.deleteUser(serverDB, userId);
+
+        const user = await serverDB.query.users.findFirst({
+          where: eq(users.id, userId),
+        });
+
+        expect(user).toBeUndefined();
+      });
+    });
+
+    describe('findById', () => {
+      it('should find user by id', async () => {
+        const user = await UserModel.findById(serverDB, userId);
+
+        expect(user).toBeDefined();
+        expect(user?.email).toBe('test@example.com');
+      });
+
+      it('should return undefined for non-existent user', async () => {
+        const user = await UserModel.findById(serverDB, 'non-existent');
+
+        expect(user).toBeUndefined();
+      });
+    });
+
+    describe('findByEmail', () => {
+      it('should find user by email', async () => {
+        const user = await UserModel.findByEmail(serverDB, 'test@example.com');
+
+        expect(user).toBeDefined();
+        expect(user?.id).toBe(userId);
+      });
+
+      it('should return undefined for non-existent email', async () => {
+        const user = await UserModel.findByEmail(serverDB, 'nonexistent@example.com');
+
+        expect(user).toBeUndefined();
+      });
+    });
+
+    describe('getUserApiKeys', () => {
+      it('should return decrypted API keys', async () => {
+        await serverDB.insert(userSettings).values({
+          id: userId,
+          keyVaults: 'encrypted-keys',
+        });
+
+        const decryptor = vi.fn().mockResolvedValue({
+          openai: 'sk-xxx',
+        });
+
+        const result = await UserModel.getUserApiKeys(serverDB, userId, decryptor);
+
+        expect(decryptor).toHaveBeenCalledWith('encrypted-keys', userId);
+        expect(result).toEqual({ openai: 'sk-xxx' });
+      });
+
+      it('should throw UserNotFoundError when settings not found', async () => {
+        await expect(
+          UserModel.getUserApiKeys(serverDB, 'non-existent', mockDecryptor),
+        ).rejects.toThrow(UserNotFoundError);
+      });
+    });
+  });
+});