centersl/lobehub
1
0
Fork 0
mirror of https://github.com/lobehub/lobehub.git synced 2026-03-26 13:19:34 +07:00
Code Issues Packages Projects Releases Wiki Activity

✅ test(database): improve test coverage for models and repositories (#10518)

Browse Source 
* 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 <noreply@anthropic.com>

* ✅ 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 <noreply@anthropic.com>

* ✅ 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 <noreply@anthropic.com>

* ✅ 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 <noreply@anthropic.com>

* ✅ 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 <noreply@anthropic.com>

* ✅ 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 <noreply@anthropic.com>

* 🐛 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Arvin Xu
2025-12-01 02:04:14 +08:00
committed by GitHub
parent c674434636
commit 28a56e96ce
7 changed files with 1557 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
CLAUDE.md
View File

@@ -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

105
packages/database/src/models/__tests__/agent.test.ts
View File

@@ -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);
});
});
});

163
packages/database/src/models/__tests__/document.test.ts
View File

@@ -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(

294
packages/database/src/models/__tests__/embedding.test.ts Normal file
View File

@@ -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);
});
});
});

261
packages/database/src/models/__tests__/oauthHandoff.test.ts Normal file
View File

@@ -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);
});
});
});

327
packages/database/src/models/__tests__/thread.test.ts Normal file
View File

@@ -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);
});
});
});

372
packages/database/src/models/__tests__/user.test.ts Normal file
View File

@@ -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);
});
});
});
});