From 28a56e96ce387320c75e454c1b0ffea2fadc7068 Mon Sep 17 00:00:00 2001 From: Arvin Xu Date: Mon, 1 Dec 2025 02:04:14 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=85=20test(database):=20improve=20test=20?= =?UTF-8?q?coverage=20for=20models=20and=20repositories=20(#10518)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * update * ✅ test(database): add ThreadModel unit tests Add comprehensive unit tests for ThreadModel covering: - create: thread creation with various parameters - query: fetch all threads for user - queryByTopicId: fetch threads by topic - findById: retrieve thread by id - update: update thread properties - delete: delete single thread - deleteAll: delete all user threads - User isolation tests for security 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * ✅ test(database): add EmbeddingModel unit tests Add comprehensive unit tests for EmbeddingModel covering: - create: create new embedding for a chunk - bulkCreate: batch create embeddings with conflict handling - delete: delete embedding by id - query: fetch all user embeddings - findById: retrieve embedding by id - countUsage: count total embeddings for user - User isolation tests for security 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * ✅ test(database): add OAuthHandoffModel unit tests Add comprehensive unit tests for OAuthHandoffModel covering: - create: create OAuth handoff with conflict handling - fetchAndConsume: fetch and delete credentials with TTL check - cleanupExpired: delete expired records (>5 min old) - exists: check credential existence without consuming - Expiration validation for 5-minute TTL 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * ✅ test(database): add UserModel unit tests Add comprehensive unit tests for UserModel covering: - getUserRegistrationDuration: calculate user registration duration - getUserState: get user state with settings and decryption - getUserSSOProviders: get linked SSO providers - getUserSettings: retrieve user settings - updateUser: update user properties - deleteSetting: delete user settings - updateSetting: create/update user settings (upsert) - updatePreference: merge and update user preferences - updateGuide: update user guide preferences Static methods: - makeSureUserExist: ensure user exists - createUser: create new user with duplicate check - deleteUser: delete user by id - findById: find user by id - findByEmail: find user by email - getUserApiKeys: get decrypted API keys 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * ✅ test(database): add missing DocumentModel tests Add tests for uncovered DocumentModel methods: - create: create new document - delete: delete document by id with user isolation - deleteAll: delete all user documents - query: query all documents with ordering - findById: find document by id with user isolation - update: update document with user isolation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * ✅ test(database): add user isolation tests for AgentModel Add user isolation security tests to ensure users cannot access or modify other users' knowledge base and file associations. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * 🐛 fix(database): fix flaky document ordering test Add 50ms delay before update to ensure timestamp difference for ordering test. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --------- Co-authored-by: Claude --- CLAUDE.md | 38 ++ .../src/models/__tests__/agent.test.ts | 105 ++++- .../src/models/__tests__/document.test.ts | 163 ++++++++ .../src/models/__tests__/embedding.test.ts | 294 ++++++++++++++ .../src/models/__tests__/oauthHandoff.test.ts | 261 ++++++++++++ .../src/models/__tests__/thread.test.ts | 327 +++++++++++++++ .../src/models/__tests__/user.test.ts | 372 ++++++++++++++++++ 7 files changed, 1557 insertions(+), 3 deletions(-) create mode 100644 packages/database/src/models/__tests__/embedding.test.ts create mode 100644 packages/database/src/models/__tests__/oauthHandoff.test.ts create mode 100644 packages/database/src/models/__tests__/thread.test.ts create mode 100644 packages/database/src/models/__tests__/user.test.ts diff --git a/CLAUDE.md b/CLAUDE.md index 06ff3c2e80..fe8a82b6c3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -55,6 +55,44 @@ see @.cursor/rules/typescript.mdc - **Dev**: Translate `locales/zh-CN/namespace.json` and `locales/en-US/namespace.json` locales file only for dev preview - DON'T run `pnpm i18n`, let CI auto handle it +## Linear Issue Management + +When working with Linear issues: + +1. **Retrieve issue details** before starting work using `mcp__linear-server__get_issue` +2. **Check for sub-issues**: If the issue has sub-issues, retrieve and review ALL sub-issues using `mcp__linear-server__list_issues` with `parentId` filter before starting work +3. **Update issue status** when completing tasks using `mcp__linear-server__update_issue` +4. **MUST add completion comment** using `mcp__linear-server__create_comment` + +### Completion Comment (REQUIRED) + +**Every time you complete an issue, you MUST add a comment summarizing the work done.** This is critical for: + +- Team visibility and knowledge sharing +- Code review context +- Future reference and debugging + +### IMPORTANT: Per-Issue Completion Rule + +**When working on multiple issues (e.g., parent issue with sub-issues), you MUST update status and add comment for EACH issue IMMEDIATELY after completing it.** Do NOT wait until all issues are done to update them in batch. + +**Workflow for EACH individual issue:** + +1. Complete the implementation for this specific issue +2. Run type check: `bun run type-check` +3. Run related tests if applicable +4. **IMMEDIATELY** update issue status to "Done": `mcp__linear-server__update_issue` +5. **IMMEDIATELY** add completion comment: `mcp__linear-server__create_comment` +6. Only then move on to the next issue + +**❌ Wrong approach:** + +- Complete Issue A → Complete Issue B → Complete Issue C → Update all statuses → Add all comments + +**✅ Correct approach:** + +- Complete Issue A → Update A status → Add A comment → Complete Issue B → Update B status → Add B comment → ... + ## Rules Index Some useful project rules are listed in @.cursor/rules/rules-index.mdc diff --git a/packages/database/src/models/__tests__/agent.test.ts b/packages/database/src/models/__tests__/agent.test.ts index 1ffebf2fe1..9c1818b067 100644 --- a/packages/database/src/models/__tests__/agent.test.ts +++ b/packages/database/src/models/__tests__/agent.test.ts @@ -20,9 +20,12 @@ import { getTestDB } from './_util'; const serverDB: LobeChatDatabase = await getTestDB(); const userId = 'agent-model-test-user-id'; +const userId2 = 'agent-model-test-user-id-2'; const agentModel = new AgentModel(serverDB, userId); +const agentModel2 = new AgentModel(serverDB, userId2); const knowledgeBase = { id: 'kb1', userId, name: 'knowledgeBase' }; +const knowledgeBase2 = { id: 'kb2', userId: userId2, name: 'knowledgeBase2' }; const fileList = [ { id: '1', @@ -42,11 +45,22 @@ const fileList = [ }, ]; +const fileList2 = [ + { + id: '3', + name: 'other.pdf', + url: 'https://a.com/other.pdf', + size: 1000, + fileType: 'application/pdf', + userId: userId2, + }, +]; + beforeEach(async () => { await serverDB.delete(users); - await serverDB.insert(users).values([{ id: userId }]); - await serverDB.insert(knowledgeBases).values(knowledgeBase); - await serverDB.insert(files).values(fileList); + await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]); + await serverDB.insert(knowledgeBases).values([knowledgeBase, knowledgeBase2]); + await serverDB.insert(files).values([...fileList, ...fileList2]); }); afterEach(async () => { @@ -226,6 +240,27 @@ describe('AgentModel', () => { expect(result).toBeUndefined(); }); + + it('should not delete another user agent knowledge base association', async () => { + const agent = await serverDB + .insert(agents) + .values({ userId }) + .returning() + .then((res) => res[0]); + await serverDB + .insert(agentsKnowledgeBases) + .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId }); + + // Try to delete with another user's model + await agentModel2.deleteAgentKnowledgeBase(agent.id, knowledgeBase.id); + + const result = await serverDB.query.agentsKnowledgeBases.findFirst({ + where: eq(agentsKnowledgeBases.agentId, agent.id), + }); + + // Should still exist + expect(result).toBeDefined(); + }); }); describe('toggleKnowledgeBase', () => { @@ -248,6 +283,28 @@ describe('AgentModel', () => { expect(result?.enabled).toBe(false); }); + + it('should not toggle another user agent knowledge base association', async () => { + const agent = await serverDB + .insert(agents) + .values({ userId }) + .returning() + .then((res) => res[0]); + + await serverDB + .insert(agentsKnowledgeBases) + .values({ agentId: agent.id, knowledgeBaseId: knowledgeBase.id, userId, enabled: true }); + + // Try to toggle with another user's model + await agentModel2.toggleKnowledgeBase(agent.id, knowledgeBase.id, false); + + const result = await serverDB.query.agentsKnowledgeBases.findFirst({ + where: eq(agentsKnowledgeBases.agentId, agent.id), + }); + + // Should still be enabled + expect(result?.enabled).toBe(true); + }); }); describe('createAgentFiles', () => { @@ -363,6 +420,26 @@ describe('AgentModel', () => { expect(result).toBeUndefined(); }); + + it('should not delete another user agent file association', async () => { + const agent = await serverDB + .insert(agents) + .values({ userId }) + .returning() + .then((res) => res[0]); + + await serverDB.insert(agentsFiles).values({ agentId: agent.id, fileId: '1', userId }); + + // Try to delete with another user's model + await agentModel2.deleteAgentFile(agent.id, '1'); + + const result = await serverDB.query.agentsFiles.findFirst({ + where: eq(agentsFiles.agentId, agent.id), + }); + + // Should still exist + expect(result).toBeDefined(); + }); }); describe('toggleFile', () => { @@ -385,5 +462,27 @@ describe('AgentModel', () => { expect(result?.enabled).toBe(false); }); + + it('should not toggle another user agent file association', async () => { + const agent = await serverDB + .insert(agents) + .values({ userId }) + .returning() + .then((res) => res[0]); + + await serverDB + .insert(agentsFiles) + .values({ agentId: agent.id, fileId: '1', userId, enabled: true }); + + // Try to toggle with another user's model + await agentModel2.toggleFile(agent.id, '1', false); + + const result = await serverDB.query.agentsFiles.findFirst({ + where: eq(agentsFiles.agentId, agent.id), + }); + + // Should still be enabled + expect(result?.enabled).toBe(true); + }); }); }); diff --git a/packages/database/src/models/__tests__/document.test.ts b/packages/database/src/models/__tests__/document.test.ts index 6e26359066..849e0dc385 100644 --- a/packages/database/src/models/__tests__/document.test.ts +++ b/packages/database/src/models/__tests__/document.test.ts @@ -54,6 +54,169 @@ const createTestDocument = async (model: DocumentModel, fModel: FileModel, conte }; describe('DocumentModel', () => { + describe('create', () => { + it('should create a new document', async () => { + const { id: fileId } = await fileModel.create({ + fileType: 'text/plain', + name: 'test.txt', + size: 100, + url: 'https://example.com/test.txt', + }); + + const file = await fileModel.findById(fileId); + if (!file) throw new Error('File not found'); + + const result = await documentModel.create({ + content: 'Test content', + fileId: file.id, + fileType: 'text/plain', + source: file.url, + sourceType: 'file', + totalCharCount: 12, + totalLineCount: 1, + }); + + expect(result).toBeDefined(); + expect(result.content).toBe('Test content'); + expect(result.fileId).toBe(file.id); + }); + }); + + describe('delete', () => { + it('should delete a document', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content'); + + await documentModel.delete(documentId); + + const deleted = await documentModel.findById(documentId); + expect(deleted).toBeUndefined(); + }); + + it('should not delete document belonging to another user', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content'); + + // Try to delete with another user's model + await documentModel2.delete(documentId); + + // Document should still exist + const stillExists = await documentModel.findById(documentId); + expect(stillExists).toBeDefined(); + }); + }); + + describe('deleteAll', () => { + it('should delete all documents for the user', async () => { + await createTestDocument(documentModel, fileModel, 'First document'); + await createTestDocument(documentModel, fileModel, 'Second document'); + await createTestDocument(documentModel2, fileModel2, 'Other user document'); + + await documentModel.deleteAll(); + + const userDocs = await documentModel.query(); + const otherUserDocs = await documentModel2.query(); + + expect(userDocs).toHaveLength(0); + expect(otherUserDocs).toHaveLength(1); + }); + }); + + describe('query', () => { + it('should return all documents for the user', async () => { + await createTestDocument(documentModel, fileModel, 'First document'); + await createTestDocument(documentModel, fileModel, 'Second document'); + + const result = await documentModel.query(); + + expect(result).toHaveLength(2); + }); + + it('should only return documents for the current user', async () => { + await createTestDocument(documentModel, fileModel, 'User 1 document'); + await createTestDocument(documentModel2, fileModel2, 'User 2 document'); + + const result = await documentModel.query(); + + expect(result).toHaveLength(1); + expect(result[0].content).toBe('User 1 document'); + }); + + it('should return documents ordered by updatedAt desc', async () => { + const { documentId: doc1Id } = await createTestDocument( + documentModel, + fileModel, + 'First document', + ); + const { documentId: doc2Id } = await createTestDocument( + documentModel, + fileModel, + 'Second document', + ); + + // Wait a bit to ensure timestamp difference + await new Promise((resolve) => setTimeout(resolve, 50)); + + // Update first document to make it more recent + await documentModel.update(doc1Id, { content: 'Updated first document' }); + + const result = await documentModel.query(); + + expect(result[0].id).toBe(doc1Id); + expect(result[1].id).toBe(doc2Id); + }); + }); + + describe('findById', () => { + it('should find document by id', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content'); + + const found = await documentModel.findById(documentId); + + expect(found).toBeDefined(); + expect(found?.id).toBe(documentId); + expect(found?.content).toBe('Test content'); + }); + + it('should return undefined for non-existent document', async () => { + const found = await documentModel.findById('non-existent-id'); + + expect(found).toBeUndefined(); + }); + + it('should not find document belonging to another user', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Test content'); + + const found = await documentModel2.findById(documentId); + + expect(found).toBeUndefined(); + }); + }); + + describe('update', () => { + it('should update a document', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content'); + + await documentModel.update(documentId, { + content: 'Updated content', + totalCharCount: 15, + }); + + const updated = await documentModel.findById(documentId); + + expect(updated?.content).toBe('Updated content'); + expect(updated?.totalCharCount).toBe(15); + }); + + it('should not update document belonging to another user', async () => { + const { documentId } = await createTestDocument(documentModel, fileModel, 'Original content'); + + await documentModel2.update(documentId, { content: 'Hacked content' }); + + const unchanged = await documentModel.findById(documentId); + + expect(unchanged?.content).toBe('Original content'); + }); + }); + describe('findByFileId', () => { it('should find document by fileId', async () => { const { documentId, file } = await createTestDocument( diff --git a/packages/database/src/models/__tests__/embedding.test.ts b/packages/database/src/models/__tests__/embedding.test.ts new file mode 100644 index 0000000000..7b6539a349 --- /dev/null +++ b/packages/database/src/models/__tests__/embedding.test.ts @@ -0,0 +1,294 @@ +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { chunks, embeddings, users } from '../../schemas'; +import { LobeChatDatabase } from '../../type'; +import { EmbeddingModel } from '../embedding'; +import { getTestDB } from './_util'; +import { designThinkingQuery } from './fixtures/embedding'; + +const userId = 'embedding-user-test'; +const otherUserId = 'other-user-test'; + +const serverDB: LobeChatDatabase = await getTestDB(); +const embeddingModel = new EmbeddingModel(serverDB, userId); + +describe('EmbeddingModel', () => { + beforeEach(async () => { + await serverDB.delete(users); + await serverDB.insert(users).values([{ id: userId }, { id: otherUserId }]); + }); + + afterEach(async () => { + await serverDB.delete(users); + }); + + describe('create', () => { + it('should create a new embedding', async () => { + // Create a chunk first + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Test chunk', userId }) + .returning(); + + const id = await embeddingModel.create({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + model: 'text-embedding-ada-002', + }); + + expect(id).toBeDefined(); + + const created = await serverDB.query.embeddings.findFirst({ + where: eq(embeddings.id, id), + }); + + expect(created).toBeDefined(); + expect(created?.chunkId).toBe(chunk.id); + expect(created?.model).toBe('text-embedding-ada-002'); + expect(created?.userId).toBe(userId); + }); + }); + + describe('bulkCreate', () => { + it('should create multiple embeddings', async () => { + // Create chunks first + const [chunk1, chunk2] = await serverDB + .insert(chunks) + .values([ + { text: 'Test chunk 1', userId }, + { text: 'Test chunk 2', userId }, + ]) + .returning(); + + await embeddingModel.bulkCreate([ + { chunkId: chunk1.id, embeddings: designThinkingQuery, model: 'text-embedding-ada-002' }, + { chunkId: chunk2.id, embeddings: designThinkingQuery, model: 'text-embedding-ada-002' }, + ]); + + const created = await serverDB.query.embeddings.findMany({ + where: eq(embeddings.userId, userId), + }); + + expect(created).toHaveLength(2); + }); + + it('should handle duplicate chunkId with onConflictDoNothing', async () => { + // Create a chunk + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Test chunk', userId }) + .returning(); + + // Create first embedding + await embeddingModel.create({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + model: 'text-embedding-ada-002', + }); + + // Try to create duplicate + await embeddingModel.bulkCreate([ + { chunkId: chunk.id, embeddings: designThinkingQuery, model: 'text-embedding-3-small' }, + ]); + + const created = await serverDB.query.embeddings.findMany({ + where: eq(embeddings.chunkId, chunk.id), + }); + + // Should still have only 1 embedding due to unique constraint + expect(created).toHaveLength(1); + expect(created[0].model).toBe('text-embedding-ada-002'); + }); + }); + + describe('delete', () => { + it('should delete an embedding', async () => { + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Test chunk', userId }) + .returning(); + + const id = await embeddingModel.create({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + model: 'text-embedding-ada-002', + }); + + await embeddingModel.delete(id); + + const deleted = await serverDB.query.embeddings.findFirst({ + where: eq(embeddings.id, id), + }); + + expect(deleted).toBeUndefined(); + }); + + it('should not delete embedding belonging to another user', async () => { + // Create chunk and embedding for other user + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Other user chunk', userId: otherUserId }) + .returning(); + + const [otherEmbedding] = await serverDB + .insert(embeddings) + .values({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + model: 'text-embedding-ada-002', + userId: otherUserId, + }) + .returning(); + + await embeddingModel.delete(otherEmbedding.id); + + const stillExists = await serverDB.query.embeddings.findFirst({ + where: eq(embeddings.id, otherEmbedding.id), + }); + + expect(stillExists).toBeDefined(); + }); + }); + + describe('query', () => { + it('should return all embeddings for the user', async () => { + const [chunk1, chunk2] = await serverDB + .insert(chunks) + .values([ + { text: 'Test chunk 1', userId }, + { text: 'Test chunk 2', userId }, + ]) + .returning(); + + await serverDB.insert(embeddings).values([ + { chunkId: chunk1.id, embeddings: designThinkingQuery, userId }, + { chunkId: chunk2.id, embeddings: designThinkingQuery, userId }, + ]); + + const result = await embeddingModel.query(); + + expect(result).toHaveLength(2); + }); + + it('should only return embeddings for the current user', async () => { + const [chunk1] = await serverDB + .insert(chunks) + .values([{ text: 'Test chunk 1', userId }]) + .returning(); + + const [chunk2] = await serverDB + .insert(chunks) + .values([{ text: 'Other user chunk', userId: otherUserId }]) + .returning(); + + await serverDB.insert(embeddings).values([ + { chunkId: chunk1.id, embeddings: designThinkingQuery, userId }, + { chunkId: chunk2.id, embeddings: designThinkingQuery, userId: otherUserId }, + ]); + + const result = await embeddingModel.query(); + + expect(result).toHaveLength(1); + expect(result[0].userId).toBe(userId); + }); + }); + + describe('findById', () => { + it('should return an embedding by id', async () => { + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Test chunk', userId }) + .returning(); + + const id = await embeddingModel.create({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + model: 'text-embedding-ada-002', + }); + + const result = await embeddingModel.findById(id); + + expect(result).toBeDefined(); + expect(result?.id).toBe(id); + expect(result?.chunkId).toBe(chunk.id); + }); + + it('should return undefined for non-existent embedding', async () => { + // Use a valid UUID format that doesn't exist + const result = await embeddingModel.findById('00000000-0000-0000-0000-000000000000'); + + expect(result).toBeUndefined(); + }); + + it('should not return embedding belonging to another user', async () => { + const [chunk] = await serverDB + .insert(chunks) + .values({ text: 'Other user chunk', userId: otherUserId }) + .returning(); + + const [otherEmbedding] = await serverDB + .insert(embeddings) + .values({ + chunkId: chunk.id, + embeddings: designThinkingQuery, + userId: otherUserId, + }) + .returning(); + + const result = await embeddingModel.findById(otherEmbedding.id); + + expect(result).toBeUndefined(); + }); + }); + + describe('countUsage', () => { + it('should return the count of embeddings for the user', async () => { + const [chunk1, chunk2, chunk3] = await serverDB + .insert(chunks) + .values([ + { text: 'Test chunk 1', userId }, + { text: 'Test chunk 2', userId }, + { text: 'Test chunk 3', userId }, + ]) + .returning(); + + await serverDB.insert(embeddings).values([ + { chunkId: chunk1.id, embeddings: designThinkingQuery, userId }, + { chunkId: chunk2.id, embeddings: designThinkingQuery, userId }, + { chunkId: chunk3.id, embeddings: designThinkingQuery, userId }, + ]); + + const count = await embeddingModel.countUsage(); + + expect(count).toBe(3); + }); + + it('should return 0 when user has no embeddings', async () => { + const count = await embeddingModel.countUsage(); + + expect(count).toBe(0); + }); + + it('should only count embeddings for the current user', async () => { + const [chunk1] = await serverDB + .insert(chunks) + .values([{ text: 'Test chunk 1', userId }]) + .returning(); + + const [chunk2] = await serverDB + .insert(chunks) + .values([{ text: 'Other user chunk', userId: otherUserId }]) + .returning(); + + await serverDB.insert(embeddings).values([ + { chunkId: chunk1.id, embeddings: designThinkingQuery, userId }, + { chunkId: chunk2.id, embeddings: designThinkingQuery, userId: otherUserId }, + ]); + + const count = await embeddingModel.countUsage(); + + expect(count).toBe(1); + }); + }); +}); diff --git a/packages/database/src/models/__tests__/oauthHandoff.test.ts b/packages/database/src/models/__tests__/oauthHandoff.test.ts new file mode 100644 index 0000000000..c3303e0110 --- /dev/null +++ b/packages/database/src/models/__tests__/oauthHandoff.test.ts @@ -0,0 +1,261 @@ +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { oauthHandoffs } from '../../schemas'; +import { LobeChatDatabase } from '../../type'; +import { OAuthHandoffModel } from '../oauthHandoff'; +import { getTestDB } from './_util'; + +const serverDB: LobeChatDatabase = await getTestDB(); +const oauthHandoffModel = new OAuthHandoffModel(serverDB); + +describe('OAuthHandoffModel', () => { + beforeEach(async () => { + await serverDB.delete(oauthHandoffs); + }); + + afterEach(async () => { + await serverDB.delete(oauthHandoffs); + vi.useRealTimers(); + }); + + describe('create', () => { + it('should create a new OAuth handoff record', async () => { + const result = await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'auth-code', state: 'state-value' }, + }); + + expect(result).toBeDefined(); + expect(result.id).toBe('handoff-1'); + expect(result.client).toBe('desktop'); + expect(result.payload).toEqual({ code: 'auth-code', state: 'state-value' }); + }); + + it('should handle conflict by doing nothing', async () => { + // Create first record + await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'original-code' }, + }); + + // Try to create duplicate - should not throw + const result = await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'new-code' }, + }); + + // Result is undefined when conflict occurs with onConflictDoNothing + expect(result).toBeUndefined(); + + // Original should remain unchanged + const record = await serverDB.query.oauthHandoffs.findFirst({ + where: eq(oauthHandoffs.id, 'handoff-1'), + }); + + expect(record?.payload).toEqual({ code: 'original-code' }); + }); + }); + + describe('fetchAndConsume', () => { + it('should fetch and delete valid credentials', async () => { + await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'auth-code', state: 'state-value' }, + }); + + const result = await oauthHandoffModel.fetchAndConsume('handoff-1', 'desktop'); + + expect(result).toBeDefined(); + expect(result?.id).toBe('handoff-1'); + expect(result?.payload).toEqual({ code: 'auth-code', state: 'state-value' }); + + // Verify it was deleted + const deleted = await serverDB.query.oauthHandoffs.findFirst({ + where: eq(oauthHandoffs.id, 'handoff-1'), + }); + expect(deleted).toBeUndefined(); + }); + + it('should return null for non-existent credentials', async () => { + const result = await oauthHandoffModel.fetchAndConsume('non-existent', 'desktop'); + + expect(result).toBeNull(); + }); + + it('should return null when client type does not match', async () => { + await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'auth-code' }, + }); + + const result = await oauthHandoffModel.fetchAndConsume('handoff-1', 'browser-extension'); + + expect(result).toBeNull(); + + // Record should still exist + const record = await serverDB.query.oauthHandoffs.findFirst({ + where: eq(oauthHandoffs.id, 'handoff-1'), + }); + expect(record).toBeDefined(); + }); + + it('should return null for expired credentials (older than 5 minutes)', async () => { + // Create a record with old timestamp + const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000); + + await serverDB.insert(oauthHandoffs).values({ + id: 'handoff-expired', + client: 'desktop', + payload: { code: 'auth-code' }, + createdAt: sixMinutesAgo, + updatedAt: sixMinutesAgo, + }); + + const result = await oauthHandoffModel.fetchAndConsume('handoff-expired', 'desktop'); + + expect(result).toBeNull(); + }); + + it('should return valid credentials within 5 minutes', async () => { + // Create a record with timestamp 4 minutes ago + const fourMinutesAgo = new Date(Date.now() - 4 * 60 * 1000); + + await serverDB.insert(oauthHandoffs).values({ + id: 'handoff-valid', + client: 'desktop', + payload: { code: 'auth-code' }, + createdAt: fourMinutesAgo, + updatedAt: fourMinutesAgo, + }); + + const result = await oauthHandoffModel.fetchAndConsume('handoff-valid', 'desktop'); + + expect(result).toBeDefined(); + expect(result?.id).toBe('handoff-valid'); + }); + }); + + describe('cleanupExpired', () => { + it('should delete expired records', async () => { + const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000); + const now = new Date(); + + // Create expired record + await serverDB.insert(oauthHandoffs).values({ + id: 'handoff-expired', + client: 'desktop', + payload: { code: 'expired-code' }, + createdAt: sixMinutesAgo, + updatedAt: sixMinutesAgo, + }); + + // Create valid record + await serverDB.insert(oauthHandoffs).values({ + id: 'handoff-valid', + client: 'desktop', + payload: { code: 'valid-code' }, + createdAt: now, + updatedAt: now, + }); + + await oauthHandoffModel.cleanupExpired(); + + // Verify expired record is deleted + const expiredRecord = await serverDB.query.oauthHandoffs.findFirst({ + where: eq(oauthHandoffs.id, 'handoff-expired'), + }); + expect(expiredRecord).toBeUndefined(); + + // Verify valid record still exists + const validRecord = await serverDB.query.oauthHandoffs.findFirst({ + where: eq(oauthHandoffs.id, 'handoff-valid'), + }); + expect(validRecord).toBeDefined(); + }); + + it('should not throw when no expired records exist', async () => { + await expect(oauthHandoffModel.cleanupExpired()).resolves.not.toThrow(); + }); + + it('should delete multiple expired records', async () => { + const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000); + + await serverDB.insert(oauthHandoffs).values([ + { + id: 'handoff-expired-1', + client: 'desktop', + payload: { code: 'code-1' }, + createdAt: sixMinutesAgo, + updatedAt: sixMinutesAgo, + }, + { + id: 'handoff-expired-2', + client: 'browser', + payload: { code: 'code-2' }, + createdAt: sixMinutesAgo, + updatedAt: sixMinutesAgo, + }, + ]); + + await oauthHandoffModel.cleanupExpired(); + + // Verify all expired records are deleted + const remainingRecords = await serverDB.query.oauthHandoffs.findMany(); + expect(remainingRecords).toHaveLength(0); + }); + }); + + describe('exists', () => { + it('should return true for existing valid credentials', async () => { + await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'auth-code' }, + }); + + const result = await oauthHandoffModel.exists('handoff-1', 'desktop'); + + expect(result).toBe(true); + }); + + it('should return false for non-existent credentials', async () => { + const result = await oauthHandoffModel.exists('non-existent', 'desktop'); + + expect(result).toBe(false); + }); + + it('should return false when client type does not match', async () => { + await oauthHandoffModel.create({ + id: 'handoff-1', + client: 'desktop', + payload: { code: 'auth-code' }, + }); + + const result = await oauthHandoffModel.exists('handoff-1', 'browser-extension'); + + expect(result).toBe(false); + }); + + it('should return false for expired credentials', async () => { + const sixMinutesAgo = new Date(Date.now() - 6 * 60 * 1000); + + await serverDB.insert(oauthHandoffs).values({ + id: 'handoff-expired', + client: 'desktop', + payload: { code: 'auth-code' }, + createdAt: sixMinutesAgo, + updatedAt: sixMinutesAgo, + }); + + const result = await oauthHandoffModel.exists('handoff-expired', 'desktop'); + + expect(result).toBe(false); + }); + }); +}); diff --git a/packages/database/src/models/__tests__/thread.test.ts b/packages/database/src/models/__tests__/thread.test.ts new file mode 100644 index 0000000000..de8a1cf4ab --- /dev/null +++ b/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); + }); + }); +}); diff --git a/packages/database/src/models/__tests__/user.test.ts b/packages/database/src/models/__tests__/user.test.ts new file mode 100644 index 0000000000..658fa58911 --- /dev/null +++ b/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); + }); + }); + }); +});