|
| 1 | +/** |
| 2 | + * Knowledge-base document totals against real PostgreSQL: the public v1 list and detail count |
| 3 | + * only the documents their caller can read, and the internal list reads no document at all |
| 4 | + * unless the caller asks for totals. |
| 5 | + */ |
| 6 | +import type { Principal } from '@sim/auth/principal' |
| 7 | +import { db } from '@sim/db' |
| 8 | +import { document, organization, user, workspace } from '@sim/db/schema' |
| 9 | +import { createMockRequest } from '@sim/testing' |
| 10 | +import { generateId } from '@sim/utils/id' |
| 11 | +import { eq, inArray } from 'drizzle-orm' |
| 12 | +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' |
| 13 | + |
| 14 | +const caller = vi.hoisted(() => ({ userId: '' })) |
| 15 | + |
| 16 | +vi.mock('@/app/api/v1/middleware', async (importOriginal) => ({ |
| 17 | + ...(await importOriginal<typeof import('@/app/api/v1/middleware')>()), |
| 18 | + authenticateRequest: async () => ({ |
| 19 | + requestId: 'fixture-request', |
| 20 | + userId: caller.userId, |
| 21 | + rateLimit: { |
| 22 | + allowed: true, |
| 23 | + remaining: 1, |
| 24 | + limit: 1, |
| 25 | + resetAt: new Date(), |
| 26 | + userId: caller.userId, |
| 27 | + keyType: 'personal', |
| 28 | + }, |
| 29 | + }), |
| 30 | +})) |
| 31 | + |
| 32 | +import { |
| 33 | + createKnowledgeAclFixtureIds, |
| 34 | + seedKnowledgeAclFixture, |
| 35 | +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' |
| 36 | +import { listInternalKnowledgeBases } from '@/lib/knowledge/application/knowledge-bases' |
| 37 | +import { GET as getKnowledgeBase } from '@/app/api/v1/knowledge/[id]/route' |
| 38 | +import { GET as listKnowledgeBases } from '@/app/api/v1/knowledge/route' |
| 39 | + |
| 40 | +const ids = createKnowledgeAclFixtureIds() |
| 41 | +const reader: Principal = { kind: 'session', userId: ids.bobId, sessionId: 'fixture-reader' } |
| 42 | + |
| 43 | +describe('knowledge-base document totals in PostgreSQL', () => { |
| 44 | + beforeAll(async () => { |
| 45 | + vi.stubGlobal('fetch', async () => { |
| 46 | + throw new Error('Unexpected provider request in knowledge-base count tests') |
| 47 | + }) |
| 48 | + caller.userId = ids.bobId |
| 49 | + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) |
| 50 | + await db.insert(document).values([ |
| 51 | + { |
| 52 | + id: generateId(), |
| 53 | + knowledgeBaseId: ids.knowledgeBaseId, |
| 54 | + filename: 'Uploaded handbook', |
| 55 | + fileUrl: 'https://fixture.test/uploaded', |
| 56 | + fileSize: 10, |
| 57 | + mimeType: 'text/plain', |
| 58 | + tokenCount: 10, |
| 59 | + processingStatus: 'completed', |
| 60 | + }, |
| 61 | + { |
| 62 | + id: generateId(), |
| 63 | + knowledgeBaseId: ids.knowledgeBaseId, |
| 64 | + filename: 'Private source document', |
| 65 | + fileUrl: 'https://fixture.test/private', |
| 66 | + fileSize: 20, |
| 67 | + mimeType: 'text/plain', |
| 68 | + tokenCount: 20, |
| 69 | + processingStatus: 'completed', |
| 70 | + connectorId: ids.connectorId, |
| 71 | + externalId: 'private-fixture', |
| 72 | + contentHash: 'fixture', |
| 73 | + acl: [`u:${ids.aliceId}@fixture.test`], |
| 74 | + aclVerifiedAt: new Date(), |
| 75 | + }, |
| 76 | + ]) |
| 77 | + }) |
| 78 | + |
| 79 | + afterAll(async () => { |
| 80 | + await db.delete(workspace).where(eq(workspace.id, ids.workspaceId)) |
| 81 | + await db.delete(organization).where(eq(organization.id, ids.organizationId)) |
| 82 | + await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId])) |
| 83 | + vi.unstubAllGlobals() |
| 84 | + }) |
| 85 | + |
| 86 | + it('lists only the documents a v1 caller can read', async () => { |
| 87 | + const response = await listKnowledgeBases( |
| 88 | + createMockRequest( |
| 89 | + 'GET', |
| 90 | + undefined, |
| 91 | + {}, |
| 92 | + `http://localhost/api/v1/knowledge?workspaceId=${ids.workspaceId}` |
| 93 | + ), |
| 94 | + { params: Promise.resolve({}) } |
| 95 | + ) |
| 96 | + expect(response.status).toBe(200) |
| 97 | + const { data } = await response.json() |
| 98 | + expect(data.knowledgeBases).toEqual([ |
| 99 | + expect.objectContaining({ id: ids.knowledgeBaseId, docCount: 1, tokenCount: 10 }), |
| 100 | + ]) |
| 101 | + }) |
| 102 | + |
| 103 | + it('details only the documents a v1 caller can read', async () => { |
| 104 | + const response = await getKnowledgeBase( |
| 105 | + createMockRequest( |
| 106 | + 'GET', |
| 107 | + undefined, |
| 108 | + {}, |
| 109 | + `http://localhost/api/v1/knowledge/${ids.knowledgeBaseId}?workspaceId=${ids.workspaceId}` |
| 110 | + ), |
| 111 | + { params: Promise.resolve({ id: ids.knowledgeBaseId }) } |
| 112 | + ) |
| 113 | + expect(response.status).toBe(200) |
| 114 | + const { data } = await response.json() |
| 115 | + expect(data.knowledgeBase).toMatchObject({ |
| 116 | + id: ids.knowledgeBaseId, |
| 117 | + docCount: 1, |
| 118 | + tokenCount: 10, |
| 119 | + }) |
| 120 | + }) |
| 121 | + |
| 122 | + it('omits totals from the internal list unless the caller asks for them', async () => { |
| 123 | + const input = { workspaceId: ids.workspaceId, scope: 'active' } as const |
| 124 | + const [plain] = (await listInternalKnowledgeBases.execute({ principal: reader, input })) |
| 125 | + .knowledgeBases |
| 126 | + expect(plain).not.toHaveProperty('docCount') |
| 127 | + expect(plain).not.toHaveProperty('tokenCount') |
| 128 | + |
| 129 | + const [counted] = ( |
| 130 | + await listInternalKnowledgeBases.execute({ |
| 131 | + principal: reader, |
| 132 | + input: { ...input, includeCounts: true }, |
| 133 | + }) |
| 134 | + ).knowledgeBases |
| 135 | + expect(counted).toMatchObject({ docCount: 1, tokenCount: 10 }) |
| 136 | + }) |
| 137 | +}) |
0 commit comments