From 4bb335b62718645f0b6035e47af39a0039aaffc8 Mon Sep 17 00:00:00 2001 From: Andre Rabold Date: Tue, 6 May 2025 08:34:16 -0700 Subject: [PATCH 1/2] refactor: normalize database schema with libraries table (fixes #86) - Added migration to introduce libraries table and update documents to use library_id foreign key - Refactored DocumentStore to use library_id internally, preserving public API - Updated all prepared statements and queries to join with libraries table - Updated and simplified tests to match new schema and mocking guidelines - Improved testing conventions in ARCHITECTURE.md to focus on public API and robust mocking - Added migration and test for schema normalization - Updated documentation to reflect new schema and testing approach --- ARCHITECTURE.md | 42 ++---- db/migrations/002-normalize-library-table.sql | 29 ++++ src/store/DocumentStore.test.ts | 131 ++++++++++++++---- src/store/DocumentStore.ts | 124 +++++++++++------ src/store/applyMigrations.test.ts | 80 +++++++++++ 5 files changed, 304 insertions(+), 102 deletions(-) create mode 100644 db/migrations/002-normalize-library-table.sql create mode 100644 src/store/applyMigrations.test.ts diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 4a2a128..3ada9b3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -490,41 +490,27 @@ This hierarchy ensures: ## Testing Conventions -This section outlines conventions and best practices for writing tests within this project. +Our testing philosophy emphasizes verifying public contracts and ensuring tests are robust and maintainable. -### Mocking with Vitest +### 1. Test Public API Behavior -When mocking modules or functions using `vitest`, it's crucial to follow a specific order due to how `vi.mock` hoisting works. `vi.mock` calls are moved to the top of the file before any imports. This means you cannot define helper functions _before_ `vi.mock` and then use them _within_ the mock setup directly. +- Focus tests on public methods: verify correct outputs or observable side effects for given inputs. +- Avoid assertions on internal implementation details (private methods, internal state). Tests should remain resilient to refactoring. -To correctly mock dependencies, follow these steps: +### 2. Mocking Principles -1. **Declare the Mock:** Call `vi.mock('./path/to/module-to-mock')` at the top of your test file, before any imports or other code. -2. **Define Mock Implementations:** _After_ the `vi.mock` call, define any helper functions, variables, or mock implementations you'll need. -3. **Import the Actual Module:** Import the specific functions or classes you intend to mock from the original module. -4. **Apply the Mock:** Use the defined mock implementations to replace the behavior of the imported functions/classes. You might need to cast the imported item as a `Mock` type (`import { type Mock } from 'vitest'`). +- Mock only true external dependencies (e.g., databases, external APIs, file system). +- When mocking modules with Vitest, be aware that `vi.mock` is hoisted. Define top-level mock functions/objects before the `vi.mock` call, and assign them within the mock factory. +- Set default behaviors for mocks globally; override them locally in tests or describe blocks as needed. +- Use shared spies for repeated calls (e.g., a database statement's `.all()`), and reset them in `beforeEach`. -**Example Structure:** +### 3. Test Structure & Assertions -```typescript -import { vi, type Mock } from "vitest"; +- Organize related tests with `describe`, and use `beforeEach`/`afterEach` for setup and cleanup. +- Assert expected return values and observable side effects. Use `expect(...).resolves/.rejects` for async code. +- Only assert direct calls to mocks if that interaction is a key part of the contract being tested. -// 1. Declare the mock (hoisted to top) -vi.mock("./dependency"); - -// 2. Define mock function/variable *after* vi.mock -const mockImplementation = vi.fn(() => "mocked result"); - -// 3. Import the actual function/class *after* defining mocks -import { functionToMock } from "./dependency"; - -// 4. Apply the mock implementation -(functionToMock as Mock).mockImplementation(mockImplementation); - -// ... rest of your test code using the mocked functionToMock ... -// expect(functionToMock()).toBe('mocked result'); -``` - -This structure ensures that mocks are set up correctly before the modules that depend on them are imported and used in your tests. +These guidelines help ensure tests are clear, maintainable, and focused on the system's observable behavior. ## Releasing diff --git a/db/migrations/002-normalize-library-table.sql b/db/migrations/002-normalize-library-table.sql new file mode 100644 index 0000000..1adc28e --- /dev/null +++ b/db/migrations/002-normalize-library-table.sql @@ -0,0 +1,29 @@ +-- Migration: Normalize schema by introducing libraries table and linking documents + +-- 1. Create libraries table +CREATE TABLE IF NOT EXISTS libraries ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL UNIQUE, + initial_url TEXT NULL, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + last_updated_at DATETIME NULL, + last_scrape_duration_ms INTEGER NULL +); + +-- 2. Add library_id to documents +ALTER TABLE documents ADD COLUMN library_id INTEGER REFERENCES libraries(id); + +-- 3. Populate libraries table from existing documents +INSERT OR IGNORE INTO libraries (name) +SELECT DISTINCT library FROM documents; + +-- 4. Update documents.library_id based on libraries table +UPDATE documents +SET library_id = ( + SELECT id FROM libraries WHERE libraries.name = documents.library +); + +-- 5. Add index on documents.library_id +CREATE INDEX IF NOT EXISTS idx_documents_library_id ON documents(library_id); + +-- Note: Handling of documents_vec and FTS triggers will be addressed in a follow-up migration. diff --git a/src/store/DocumentStore.test.ts b/src/store/DocumentStore.test.ts index 2899a42..f7a2ae8 100644 --- a/src/store/DocumentStore.test.ts +++ b/src/store/DocumentStore.test.ts @@ -1,4 +1,13 @@ -import { type Mock, afterAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { + type Mock, + afterAll, + afterEach, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; import { VECTOR_DIMENSION } from "./types"; // --- Mocking Setup --- @@ -16,17 +25,19 @@ import { createEmbeddingModel } from "./embeddings/EmbeddingFactory"; embedDocuments: vi.fn(), }); -// Mock better-sqlite3 +/** + * Initial generic mocks for better-sqlite3. + * Will be replaced with dynamic mocks after vi.mock due to hoisting. + */ const mockStatementAll = vi.fn().mockReturnValue([]); -// Ensure the mock statement object covers methods used by *all* statements prepared in DocumentStore const mockStatement = { all: mockStatementAll, - run: vi.fn().mockReturnValue({ changes: 0, lastInsertRowid: 1 }), // Mock run for insert/delete - get: vi.fn().mockReturnValue(undefined), // Mock get for getById/checkExists etc. + run: vi.fn().mockReturnValue({ changes: 0, lastInsertRowid: 1 }), + get: vi.fn().mockReturnValue(undefined), }; -const mockPrepare = vi.fn().mockReturnValue(mockStatement); +let mockPrepare = vi.fn().mockReturnValue(mockStatement); const mockDb = { - prepare: mockPrepare, + prepare: (...args: unknown[]) => mockPrepare(...args), exec: vi.fn(), transaction: vi.fn( (fn) => @@ -36,9 +47,20 @@ const mockDb = { close: vi.fn(), }; vi.mock("better-sqlite3", () => ({ - default: vi.fn(() => mockDb), // Mock the default export (constructor) + default: vi.fn(() => mockDb), })); +/** + * Simplified mockPrepare: always returns a generic statement object. + * Test-specific SQL overrides are set up in each test/describe as needed. + */ +mockPrepare = vi.fn(() => ({ + get: vi.fn().mockReturnValue(undefined), + run: vi.fn().mockReturnValue({ changes: 1, lastInsertRowid: 1 }), + all: mockStatementAll, +})); +mockDb.prepare = (...args: unknown[]) => mockPrepare(...args); + // Mock sqlite-vec vi.mock("sqlite-vec", () => ({ load: vi.fn(), @@ -59,14 +81,14 @@ describe("DocumentStore", () => { beforeEach(async () => { vi.clearAllMocks(); // Clear call history etc. + mockStatementAll.mockClear(); + mockStatementAll.mockReturnValue([]); // Reset the mock factory implementation for this test run (createEmbeddingModel as ReturnType).mockReturnValue({ embedQuery: mockEmbedQuery, embedDocuments: mockEmbedDocuments, }); - mockPrepare.mockReturnValue(mockStatement); // <-- Re-configure prepare mock return value - // Reset embedQuery to handle initialization vector mockEmbedQuery.mockResolvedValue(new Array(VECTOR_DIMENSION).fill(0.1)); @@ -284,35 +306,85 @@ describe("DocumentStore", () => { }); describe("Embedding Model Dimensions", () => { + let getLibraryIdByNameMock: Mock; + let insertLibraryMock: Mock; + let lastInsertedVector: number[]; + + beforeEach(() => { + getLibraryIdByNameMock = vi.fn().mockReturnValue({ id: 1 }); + insertLibraryMock = vi.fn().mockReturnValue({ changes: 1, lastInsertRowid: 1 }); + lastInsertedVector = []; + + mockPrepare.mockImplementation((sql: string) => { + if (sql.includes("SELECT id FROM libraries WHERE name = ?")) { + return { + get: getLibraryIdByNameMock, + run: vi.fn(), + all: mockStatementAll, + }; + } + if (sql.includes("INSERT INTO libraries")) { + return { + run: insertLibraryMock, + get: vi.fn(), + all: mockStatementAll, + }; + } + if (sql.includes("INSERT INTO documents_vec")) { + return { + run: vi.fn((...args) => { + if (typeof args[3] === "string") { + try { + const arr = JSON.parse(args[3]); + if (Array.isArray(arr)) lastInsertedVector = arr; + } catch {} + } + return { changes: 1, lastInsertRowid: 1 }; + }), + get: vi.fn(), + all: mockStatementAll, + }; + } + return { + get: vi.fn(), + run: vi.fn().mockReturnValue({ changes: 1, lastInsertRowid: 1 }), + all: mockStatementAll, + }; + }); + }); + + afterEach(() => { + mockPrepare.mockImplementation(() => ({ + get: vi.fn().mockReturnValue(undefined), + run: vi.fn().mockReturnValue({ changes: 1, lastInsertRowid: 1 }), + all: mockStatementAll, + })); + }); + it("should accept a model that produces ${VECTOR_DIMENSION}-dimensional vectors", async () => { - // Mock a ${VECTOR_DIMENSION}-dimensional vector mockEmbedQuery.mockResolvedValueOnce(new Array(VECTOR_DIMENSION).fill(0.1)); documentStore = new DocumentStore(":memory:"); await expect(documentStore.initialize()).resolves.not.toThrow(); }); it("should accept and pad vectors from models with smaller dimensions", async () => { - // Mock 768-dimensional vectors mockEmbedQuery.mockResolvedValueOnce(new Array(768).fill(0.1)); mockEmbedDocuments.mockResolvedValueOnce([new Array(768).fill(0.1)]); documentStore = new DocumentStore(":memory:"); await documentStore.initialize(); - // Should pad to ${VECTOR_DIMENSION} when inserting const doc = { pageContent: "test content", metadata: { title: "test", url: "http://test.com", path: ["test"] }, }; - // This should succeed (vectors are padded internally) await expect( documentStore.addDocuments("test-lib", "1.0.0", [doc]), ).resolves.not.toThrow(); }); it("should reject models that produce vectors larger than ${VECTOR_DIMENSION} dimensions", async () => { - // Mock a 3072-dimensional vector (like text-embedding-3-large) mockEmbedQuery.mockResolvedValueOnce(new Array(3072).fill(0.1)); documentStore = new DocumentStore(":memory:"); await expect(documentStore.initialize()).rejects.toThrow( @@ -321,12 +393,11 @@ describe("DocumentStore", () => { }); it("should pad both document and query vectors consistently", async () => { - // Mock 768-dimensional vectors for both init and subsequent operations const smallVector = new Array(768).fill(0.1); mockEmbedQuery - .mockResolvedValueOnce(smallVector) // for initialization - .mockResolvedValueOnce(smallVector); // for search query - mockEmbedDocuments.mockResolvedValueOnce([smallVector]); // for document embeddings + .mockResolvedValueOnce(smallVector) + .mockResolvedValueOnce(smallVector); + mockEmbedDocuments.mockResolvedValueOnce([smallVector]); documentStore = new DocumentStore(":memory:"); await documentStore.initialize(); @@ -336,24 +407,26 @@ describe("DocumentStore", () => { metadata: { title: "test", url: "http://test.com", path: ["test"] }, }; - // Add a document (this pads the document vector) + mockStatementAll.mockImplementationOnce(() => [ + { + id: "id1", + content: "content", + metadata: JSON.stringify({}), + vec_score: 1, + fts_score: 1, + }, + ]); + await documentStore.addDocuments("test-lib", "1.0.0", [doc]); - // Search should work (query vector gets padded too) await expect( documentStore.findByContent("test-lib", "1.0.0", "test query", 5), ).resolves.not.toThrow(); - // Verify both vectors were padded (via the JSON stringification) - const insertCall = mockStatement.run.mock.calls.find( - (call) => call[0]?.toString().startsWith("1"), // Looking for rowid=1 - ); const searchCall = mockStatementAll.mock.lastCall; - - // Both vectors should be stringified arrays of length ${VECTOR_DIMENSION} - const insertVector = JSON.parse(insertCall?.[3] || "[]"); const searchVector = JSON.parse(searchCall?.[2] || "[]"); - expect(insertVector.length).toBe(VECTOR_DIMENSION); + + expect(lastInsertedVector.length).toBe(VECTOR_DIMENSION); expect(searchVector.length).toBe(VECTOR_DIMENSION); }); }); diff --git a/src/store/DocumentStore.ts b/src/store/DocumentStore.ts index acfb0fe..5391cd8 100644 --- a/src/store/DocumentStore.ts +++ b/src/store/DocumentStore.ts @@ -49,6 +49,8 @@ export class DocumentStore { getPrecedingSiblings: Database.Statement; getSubsequentSiblings: Database.Statement; getParentChunk: Database.Statement; + insertLibrary: Database.Statement; + getLibraryIdByName: Database.Statement; }; /** @@ -117,19 +119,32 @@ export class DocumentStore { const statements = { getById: this.db.prepare("SELECT * FROM documents WHERE id = ?"), insertDocument: this.db.prepare( - "INSERT INTO documents (library, version, url, content, metadata, sort_order, indexed_at) VALUES (?, ?, ?, ?, ?, ?, ?)", // Added indexed_at + "INSERT INTO documents (library_id, version, url, content, metadata, sort_order, indexed_at) VALUES (?, ?, ?, ?, ?, ?, ?)", ), insertEmbedding: this.db.prepare<[number, string]>( "INSERT INTO documents_vec (rowid, library, version, embedding) VALUES (?, ?, ?, ?)", ), + insertLibrary: this.db.prepare( + "INSERT INTO libraries (name) VALUES (?) ON CONFLICT(name) DO NOTHING", + ), + getLibraryIdByName: this.db.prepare("SELECT id FROM libraries WHERE name = ?"), deleteDocuments: this.db.prepare( - "DELETE FROM documents WHERE library = ? AND version = ?", + `DELETE FROM documents + WHERE library_id = (SELECT id FROM libraries WHERE name = ?) + AND version = ?`, ), queryVersions: this.db.prepare( - "SELECT DISTINCT version FROM documents WHERE library = ? ORDER BY version", + `SELECT DISTINCT d.version + FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + ORDER BY d.version`, ), checkExists: this.db.prepare( - "SELECT id FROM documents WHERE library = ? AND version = ? LIMIT 1", + `SELECT id FROM documents + WHERE library_id = (SELECT id FROM libraries WHERE name = ?) + AND version = ? + LIMIT 1`, ), queryLibraryVersions: this.db.prepare( `SELECT @@ -143,44 +158,48 @@ export class DocumentStore { ORDER BY library, version`, ), getChildChunks: this.db.prepare(` - SELECT * FROM documents - WHERE library = ? - AND version = ? - AND url = ? - AND json_array_length(json_extract(metadata, '$.path')) = ? - AND json_extract(metadata, '$.path') LIKE ? || '%' - AND sort_order > (SELECT sort_order FROM documents WHERE id = ?) - ORDER BY sort_order + SELECT d.* FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + AND d.version = ? + AND d.url = ? + AND json_array_length(json_extract(d.metadata, '$.path')) = ? + AND json_extract(d.metadata, '$.path') LIKE ? || '%' + AND d.sort_order > (SELECT sort_order FROM documents WHERE id = ?) + ORDER BY d.sort_order LIMIT ? `), getPrecedingSiblings: this.db.prepare(` - SELECT * FROM documents - WHERE library = ? - AND version = ? - AND url = ? - AND sort_order < (SELECT sort_order FROM documents WHERE id = ?) - AND json_extract(metadata, '$.path') = ? - ORDER BY sort_order DESC + SELECT d.* FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + AND d.version = ? + AND d.url = ? + AND d.sort_order < (SELECT sort_order FROM documents WHERE id = ?) + AND json_extract(d.metadata, '$.path') = ? + ORDER BY d.sort_order DESC LIMIT ? `), getSubsequentSiblings: this.db.prepare(` - SELECT * FROM documents - WHERE library = ? - AND version = ? - AND url = ? - AND sort_order > (SELECT sort_order FROM documents WHERE id = ?) - AND json_extract(metadata, '$.path') = ? - ORDER BY sort_order + SELECT d.* FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + AND d.version = ? + AND d.url = ? + AND d.sort_order > (SELECT sort_order FROM documents WHERE id = ?) + AND json_extract(d.metadata, '$.path') = ? + ORDER BY d.sort_order LIMIT ? `), getParentChunk: this.db.prepare(` - SELECT * FROM documents - WHERE library = ? - AND version = ? - AND url = ? - AND json_extract(metadata, '$.path') = ? - AND sort_order < (SELECT sort_order FROM documents WHERE id = ?) - ORDER BY sort_order DESC + SELECT d.* FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + AND d.version = ? + AND d.url = ? + AND json_extract(d.metadata, '$.path') = ? + AND d.sort_order < (SELECT sort_order FROM documents WHERE id = ?) + ORDER BY d.sort_order DESC LIMIT 1 `), }; @@ -387,6 +406,16 @@ export class DocumentStore { } const paddedEmbeddings = rawEmbeddings.map((vector) => this.padVector(vector)); + // Insert or get library_id + this.statements.insertLibrary.run(library.toLowerCase()); + const libraryIdRow = this.statements.getLibraryIdByName.get( + library.toLowerCase(), + ) as { id: number } | undefined; + if (!libraryIdRow || typeof libraryIdRow.id !== "number") { + throw new StoreError(`Failed to resolve library_id for library: ${library}`); + } + const libraryId = libraryIdRow.id; + // Insert documents in a transaction const transaction = this.db.transaction((docs: typeof documents) => { for (let i = 0; i < docs.length; i++) { @@ -398,7 +427,7 @@ export class DocumentStore { // Insert into main documents table const result = this.statements.insertDocument.run( - library.toLowerCase(), + libraryId, version.toLowerCase(), url, doc.pageContent, @@ -408,7 +437,7 @@ export class DocumentStore { ); const rowId = result.lastInsertRowid; - // Insert into vector table + // Insert into vector table (still uses library/version for now) this.statements.insertEmbedding.run( BigInt(rowId), library.toLowerCase(), @@ -476,14 +505,16 @@ export class DocumentStore { const stmt = this.db.prepare(` WITH vec_scores AS ( SELECT - rowid as id, - distance as vec_score - FROM documents_vec - WHERE library = ? - AND version = ? - AND embedding MATCH ? - ORDER BY vec_score - LIMIT ? + dv.rowid as id, + dv.distance as vec_score + FROM documents_vec dv + JOIN documents d ON dv.rowid = d.id + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? + AND d.version = ? + AND dv.embedding MATCH ? + AND dv.k = ? + ORDER BY dv.distance ), fts_scores AS ( SELECT @@ -491,7 +522,8 @@ export class DocumentStore { bm25(documents_fts, 10.0, 1.0, 5.0, 1.0) as fts_score FROM documents_fts f JOIN documents d ON f.rowid = d.rowid - WHERE d.library = ? + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? AND d.version = ? AND documents_fts MATCH ? ORDER BY fts_score @@ -703,7 +735,9 @@ export class DocumentStore { // Use parameterized query for variable number of IDs const placeholders = ids.map(() => "?").join(","); const stmt = this.db.prepare( - `SELECT * FROM documents WHERE library = ? AND version = ? AND id IN (${placeholders}) ORDER BY sort_order`, + `SELECT d.* FROM documents d + JOIN libraries l ON d.library_id = l.id + WHERE l.name = ? AND d.version = ? AND d.id IN (${placeholders}) ORDER BY d.sort_order`, ); const rows = stmt.all( library.toLowerCase(), diff --git a/src/store/applyMigrations.test.ts b/src/store/applyMigrations.test.ts new file mode 100644 index 0000000..e60ceca --- /dev/null +++ b/src/store/applyMigrations.test.ts @@ -0,0 +1,80 @@ +// Integration test for database migrations using a real SQLite database + +import Database, { type Database as DatabaseType } from "better-sqlite3"; +import * as sqliteVec from "sqlite-vec"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { applyMigrations } from "./applyMigrations"; + +describe("Database Migrations", () => { + let db: DatabaseType; + + beforeEach(() => { + db = new Database(":memory:"); + sqliteVec.load(db); + }); + + afterEach(() => { + db.close(); + }); + + it("should apply all migrations and create expected tables and columns", () => { + expect(() => applyMigrations(db)).not.toThrow(); + + // Check tables + const tables = db + .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;") + .all(); + interface TableRow { + name: string; + } + const tableNames = (tables as TableRow[]).map((t) => t.name); + expect(tableNames).toContain("documents"); + expect(tableNames).toContain("documents_fts"); + expect(tableNames).toContain("documents_vec"); + expect(tableNames).toContain("libraries"); + + // Check columns for 'documents' + const documentsColumns = db.prepare("PRAGMA table_info(documents);").all(); + interface ColumnInfo { + name: string; + } + const documentsColumnNames = (documentsColumns as ColumnInfo[]).map( + (col) => col.name, + ); + expect(documentsColumnNames).toEqual( + expect.arrayContaining([ + "id", + "library_id", + "version", + "url", + "content", + "metadata", + "sort_order", + "indexed_at", + ]), + ); + + // Check columns for 'libraries' + const librariesColumns = db.prepare("PRAGMA table_info(libraries);").all(); + const librariesColumnNames = (librariesColumns as ColumnInfo[]).map( + (col) => col.name, + ); + expect(librariesColumnNames).toEqual(expect.arrayContaining(["id", "name"])); + + // Check FTS virtual table + const ftsTableInfo = db + .prepare( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='documents_fts';", + ) + .get() as { sql: string } | undefined; + expect(ftsTableInfo?.sql).toContain("VIRTUAL TABLE documents_fts USING fts5"); + + // Check vector virtual table + const vecTableInfo = db + .prepare( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='documents_vec';", + ) + .get() as { sql: string } | undefined; + expect(vecTableInfo?.sql).toMatch(/CREATE VIRTUAL TABLE documents_vec USING vec0/i); + }); +}); From 1e24cf22a934e376e638be8b46ceae0e594b17ec Mon Sep 17 00:00:00 2001 From: Andre Rabold Date: Sat, 17 May 2025 07:37:33 -0700 Subject: [PATCH 2/2] test(store): fix batching test for addDocuments to mock library id resolution Ensure the batching test for addDocuments correctly mocks the library id lookup and re-instantiates DocumentStore after patching the mock, preventing StoreError for test-lib-large-batch. This makes the test robust and prevents false negatives when batching embeddings. --- src/store/DocumentStore.test.ts | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/src/store/DocumentStore.test.ts b/src/store/DocumentStore.test.ts index f7a2ae8..3d8b299 100644 --- a/src/store/DocumentStore.test.ts +++ b/src/store/DocumentStore.test.ts @@ -285,13 +285,31 @@ describe("DocumentStore", () => { .mockResolvedValueOnce(firstBatchEmbeddings) .mockResolvedValueOnce(secondBatchEmbeddings); - // Mock insertDocument to return sequential rowids - for (let i = 0; i < numDocuments; i++) { - mockStatement.run.mockReturnValueOnce({ - changes: 1, - lastInsertRowid: BigInt(i + 1), - }); - } + // Patch mockPrepare for this test to handle library id resolution for test-lib-large-batch + const originalMockPrepare = mockPrepare; + mockPrepare = vi.fn((sql: string) => { + if (sql.includes("SELECT id FROM libraries WHERE name = ?")) { + return { + get: (name: string) => + name === "test-lib-large-batch" ? { id: 1 } : undefined, + run: vi.fn(), + all: mockStatementAll, + }; + } + if (sql.includes("INSERT INTO libraries")) { + return { + run: vi.fn(), + get: vi.fn(), + all: mockStatementAll, + }; + } + return originalMockPrepare(sql); + }); + mockDb.prepare = (...args: unknown[]) => mockPrepare(...args); + + // Re-instantiate DocumentStore after patching mockPrepare + documentStore = new DocumentStore(":memory:"); + await documentStore.initialize(); await documentStore.addDocuments("test-lib-large-batch", "1.0.0", documents);