Skip to content

Commit

Permalink
update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Patrick-Erichsen committed Aug 15, 2024
1 parent 978f387 commit 8cdb68c
Show file tree
Hide file tree
Showing 13 changed files with 186 additions and 231 deletions.
52 changes: 22 additions & 30 deletions core/indexing/CodeSnippetsIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import { DatabaseConnection, SqliteDb, tagToString } from "./refreshIndex.js";
import {
IndexResultType,
MarkCompleteCallback,
PathAndCacheKey,
RefreshIndexResults,
type CodebaseIndex,
} from "./types.js";
Expand Down Expand Up @@ -182,47 +183,38 @@ export class CodeSnippetsCodebaseIndex implements CodebaseIndex {
// Delete
for (let i = 0; i < results.del.length; i++) {
const del = results.del[i];
const deleted = await db.run(
"DELETE FROM code_snippets WHERE path = ? AND cacheKey = ?",

const snippetToDelete = await db.get(
"SELECT id FROM code_snippets WHERE path = ? AND cacheKey = ?",
[del.path, del.cacheKey],
);

await db.run("DELETE FROM code_snippets WHERE id = ?", [
snippetToDelete.id,
]);

await db.run("DELETE FROM code_snippets_tags WHERE snippetId = ?", [
deleted.lastID,
snippetToDelete.id,
]);

markComplete([del], IndexResultType.Delete);
}

// Add tag
for (let i = 0; i < results.addTag.length; i++) {
const addTag = results.addTag[i];
let snippets: (ChunkWithoutID & { title: string })[] = [];
try {
snippets = await this.getSnippetsInFile(
addTag.path,
await this.ide.readFile(addTag.path),
);
} catch (e) {
// If can't parse, assume malformatted code
console.error(`Error parsing ${addTag.path}:`, e);
}

for (const snippet of snippets) {
const { lastID } = await db.run(
"REPLACE INTO code_snippets (path, cacheKey, content, title, startLine, endLine) VALUES (?, ?, ?, ?, ?, ?)",
[
addTag.path,
addTag.cacheKey,
snippet.content,
snippet.title,
snippet.startLine,
snippet.endLine,
],
);
await db.run(
"REPLACE INTO code_snippets_tags (snippetId, tag) VALUES (?, ?)",
[lastID, tagString],
);
}
await db.run(
`
REPLACE INTO code_snippets_tags (tag, snippetId)
SELECT ?, (
SELECT id
FROM code_snippets
WHERE cacheKey = ? AND path = ?
)
`,
[tagString, addTag.cacheKey, addTag.path],
);

markComplete([results.addTag[i]], IndexResultType.AddTag);
}
Expand Down
11 changes: 6 additions & 5 deletions core/indexing/FullTextSearch.ts
Original file line number Diff line number Diff line change
Expand Up @@ -91,11 +91,12 @@ export class FullTextSearchCodebaseIndex implements CodebaseIndex {

// Delete
for (const item of results.del) {
const { lastID } = await db.run(
"DELETE FROM fts_metadata WHERE path = ? AND cacheKey = ?",
[item.path, item.cacheKey],
);
await db.run("DELETE FROM fts WHERE rowid = ?", [lastID]);
await db.run("DELETE FROM fts_metadata WHERE path = ? AND cacheKey = ?", [
item.path,
item.cacheKey,
]);

await db.run("DELETE FROM fts WHERE path = ?", [item.path]);

markComplete([item], IndexResultType.Delete);
}
Expand Down
66 changes: 34 additions & 32 deletions core/indexing/LanceDbIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import {
PathAndCacheKey,
RefreshIndexResults,
} from "./types.js";
import lance from "vectordb";

// LanceDB converts to lowercase, so names must all be lowercase
interface LanceDbRow {
Expand Down Expand Up @@ -48,7 +49,7 @@ export class LanceDbIndex implements CodebaseIndex {
private readonly continueServerClient?: IContinueServerClient,
) {}

private tableNameForTag(tag: IndexTag) {
tableNameForTag(tag: IndexTag) {
return tagToString(tag).replace(/[^\w-_.]/g, "");
}

Expand Down Expand Up @@ -203,36 +204,35 @@ export class LanceDbIndex implements CodebaseIndex {
markComplete: MarkCompleteCallback,
repoName: string | undefined,
): AsyncGenerator<IndexingProgressUpdate> {
const lancedb = await import("vectordb");
const tableName = this.tableNameForTag(tag);
const db = await lancedb.connect(getLanceDbPath());
const sqliteDb = await SqliteDb.get();
await this.createSqliteCacheTable(sqliteDb);

const sqlite = await SqliteDb.get();
await this.createSqliteCacheTable(sqlite);
const lanceTableName = this.tableNameForTag(tag);
const lanceDb = await lance.connect(getLanceDbPath());
const existingLanceTables = await lanceDb.tableNames();

// Compute
let table: Table<number[]> | undefined = undefined;
const existingTables = await db.tableNames();
let needToCreateTable = !existingTables.includes(tableName);
let lanceTable: Table<number[]> | undefined = undefined;
let needToCreateLanceTable = !existingLanceTables.includes(lanceTableName);

// Compute
const addComputedLanceDbRows = async (
pathAndCacheKey: PathAndCacheKey,
computedRows: LanceDbRow[],
) => {
// Create table if needed, add computed rows
if (table) {
if (lanceTable) {
if (computedRows.length > 0) {
await table.add(computedRows);
await lanceTable.add(computedRows);
}
} else if (existingTables.includes(tableName)) {
table = await db.openTable(tableName);
needToCreateTable = false;
} else if (existingLanceTables.includes(lanceTableName)) {
lanceTable = await lanceDb.openTable(lanceTableName);
needToCreateLanceTable = false;
if (computedRows.length > 0) {
await table.add(computedRows);
await lanceTable.add(computedRows);
}
} else if (computedRows.length > 0) {
table = await db.createTable(tableName, computedRows);
needToCreateTable = false;
lanceTable = await lanceDb.createTable(lanceTableName, computedRows);
needToCreateLanceTable = false;
}

// Mark item complete
Expand Down Expand Up @@ -272,7 +272,7 @@ export class LanceDbIndex implements CodebaseIndex {
};
rows.push(row);

await sqlite.run(
await sqliteDb.run(
"INSERT INTO lance_db_cache (uuid, cacheKey, path, artifact_id, vector, startLine, endLine, contents) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
row.uuid,
row.cachekey,
Expand Down Expand Up @@ -306,13 +306,13 @@ export class LanceDbIndex implements CodebaseIndex {
};

const dbRows = await this.computeRows(results.compute);
await this.insertRows(sqlite, dbRows);
await this.insertRows(sqliteDb, dbRows);
await markComplete(results.compute, IndexResultType.Compute);
let accumulatedProgress = 0;

// Add tag - retrieve the computed info from lance sqlite cache
for (const { path, cacheKey } of results.addTag) {
const stmt = await sqlite.prepare(
const stmt = await sqliteDb.prepare(
"SELECT * FROM lance_db_cache WHERE cacheKey = ? AND path = ? AND artifact_id = ?",
cacheKey,
path,
Expand All @@ -330,15 +330,15 @@ export class LanceDbIndex implements CodebaseIndex {
});

if (lanceRows.length > 0) {
if (needToCreateTable) {
table = await db.createTable(tableName, lanceRows);
needToCreateTable = false;
} else if (!table) {
table = await db.openTable(tableName);
needToCreateTable = false;
await table.add(lanceRows);
if (needToCreateLanceTable) {
lanceTable = await lanceDb.createTable(lanceTableName, lanceRows);
needToCreateLanceTable = false;
} else if (!lanceTable) {
lanceTable = await lanceDb.openTable(lanceTableName);
needToCreateLanceTable = false;
await lanceTable.add(lanceRows);
} else {
await table?.add(lanceRows);
await lanceTable?.add(lanceRows);
}
}

Expand All @@ -352,11 +352,13 @@ export class LanceDbIndex implements CodebaseIndex {
}

// Delete or remove tag - remove from lance table)
if (!needToCreateTable) {
if (!needToCreateLanceTable) {
const toDel = [...results.removeTag, ...results.del];
for (const { path, cacheKey } of toDel) {
// This is where the aforementioned lowercase conversion problem shows
await table?.delete(`cachekey = '${cacheKey}' AND path = '${path}'`);
await lanceTable?.delete(
`cachekey = '${cacheKey}' AND path = '${path}'`,
);

accumulatedProgress += 1 / toDel.length / 3;
yield {
Expand All @@ -370,7 +372,7 @@ export class LanceDbIndex implements CodebaseIndex {

// Delete - also remove from sqlite cache
for (const { path, cacheKey } of results.del) {
await sqlite.run(
await sqliteDb.run(
"DELETE FROM lance_db_cache WHERE cacheKey = ? AND path = ? AND artifact_id = ?",
cacheKey,
path,
Expand Down
11 changes: 7 additions & 4 deletions core/indexing/chunk/ChunkCodebaseIndex.ts
Original file line number Diff line number Diff line change
Expand Up @@ -117,13 +117,16 @@ export class ChunkCodebaseIndex implements CodebaseIndex {

// Delete
for (const item of results.del) {
const deleted = await db.run("DELETE FROM chunks WHERE cacheKey = ?", [
item.cacheKey,
]);
const chunkToDelete = await db.get(
"SELECT id FROM chunks WHERE cacheKey = ?",
[item.cacheKey],
);

await db.run("DELETE FROM chunks WHERE id = ?", [chunkToDelete.id]);

// Delete from chunk_tags
await db.run("DELETE FROM chunk_tags WHERE chunkId = ?", [
deleted.lastID,
chunkToDelete.id,
]);

await markComplete([item], IndexResultType.Delete);
Expand Down
1 change: 1 addition & 0 deletions core/jest.config.js
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
/* eslint-disable @typescript-eslint/naming-convention */
import path from "path";
import { fileURLToPath } from "url";

Expand Down
3 changes: 0 additions & 3 deletions core/jest.global-setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,4 @@ import path from "path";

export default async function () {
process.env.CONTINUE_GLOBAL_DIR = path.join(__dirname, ".continue-test");
if (fs.existsSync(process.env.CONTINUE_GLOBAL_DIR)) {
fs.rmSync(process.env.CONTINUE_GLOBAL_DIR, { recursive: true });
}
}
15 changes: 15 additions & 0 deletions core/jest.setup-after-env.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import { TextEncoder, TextDecoder } from "util";
import fetch, { Request, Response } from "node-fetch";
import { jest } from "@jest/globals";
import fs from "fs";

if (process.env.DEBUG === "jest") {
jest.setTimeout(5 * 60 * 1000);
Expand All @@ -14,3 +15,17 @@ globalThis.Request = Request;
globalThis.Response = Response;
globalThis.TextEncoder = TextEncoder;
globalThis.TextDecoder = TextDecoder;

global.beforeAll(() => {
// We can ignore TS warnings here since we set this explicitly in `./jest.global-setup.ts`
if (fs.existsSync(process.env.CONTINUE_GLOBAL_DIR!)) {
fs.rmSync(process.env.CONTINUE_GLOBAL_DIR!, { recursive: true });
}
});

global.afterAll(() => {
// We can ignore TS warnings here since we set this explicitly in `./jest.global-setup.ts`
if (fs.existsSync(process.env.CONTINUE_GLOBAL_DIR!)) {
fs.rmSync(process.env.CONTINUE_GLOBAL_DIR!, { recursive: true });
}
});
2 changes: 1 addition & 1 deletion core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
"version": "1.0.13",
"description": "The Continue Core contains functionality that can be shared across web, VS Code, or Node.js",
"scripts": {
"test": "cross-env NODE_OPTIONS=--experimental-vm-modules jest",
"test": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --runInBand",
"build:npm": "tsc -p ./tsconfig.npm.json",
"lint": "eslint . --ext ts",
"lint:fix": "eslint . --ext ts --fix"
Expand Down
53 changes: 20 additions & 33 deletions core/test/indexing/ChunkCodebaseIndex.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,14 @@ describe("ChunkCodebaseIndex", () => {
let index: ChunkCodebaseIndex;
let db: DatabaseConnection;

async function getAllChunks() {
return await db.all("SELECT * FROM chunks");
}

async function getAllChunkTags() {
return await db.all("SELECT * FROM chunk_tags");
}

beforeAll(async () => {
const pathSep = await testIde.pathSep();

Expand All @@ -40,54 +48,33 @@ describe("ChunkCodebaseIndex", () => {

// Compute test
await updateIndexAndAwaitGenerator(index, "compute", mockMarkComplete);

const computeResult = await db.get(
"SELECT * FROM chunks WHERE cacheKey = ?",
[mockPathAndCacheKey.cacheKey],
);

expect(computeResult).toBeTruthy();
expect((await getAllChunks()).length).toBe(1);
expect((await getAllChunkTags()).length).toBe(1);
expect(mockMarkComplete).toHaveBeenCalledWith(
[mockPathAndCacheKey],
IndexResultType.Compute,
);

// AddTag test
await updateIndexAndAwaitGenerator(index, "addTag", mockMarkComplete);

const addTagResult = await db.get(
"SELECT * FROM chunk_tags WHERE chunkId = ?",
[computeResult.id],
);

expect(addTagResult).toBeTruthy();
expect(mockMarkComplete).toHaveBeenCalledWith(
[mockPathAndCacheKey],
IndexResultType.AddTag,
);

// RemoveTag test
await updateIndexAndAwaitGenerator(index, "removeTag", mockMarkComplete);

const removeTagResult = await db.get(
"SELECT * FROM chunk_tags WHERE id = ?",
[addTagResult.id],
expect((await getAllChunkTags()).length).toBe(0);
expect(mockMarkComplete).toHaveBeenCalledWith(
[mockPathAndCacheKey],
IndexResultType.RemoveTag,
);

expect(removeTagResult).toBeFalsy();
// AddTag test
await updateIndexAndAwaitGenerator(index, "addTag", mockMarkComplete);
expect((await getAllChunkTags()).length).toBe(1);
expect(mockMarkComplete).toHaveBeenCalledWith(
[mockPathAndCacheKey],
IndexResultType.RemoveTag,
IndexResultType.AddTag,
);

// Delete test
await updateIndexAndAwaitGenerator(index, "del", mockMarkComplete);

const delResult = await db.get("SELECT * FROM chunks WHERE id = ?", [
computeResult.id,
]);

expect(delResult).toBeFalsy();
expect((await getAllChunks()).length).toBe(0);
expect((await getAllChunkTags()).length).toBe(0);
expect(mockMarkComplete).toHaveBeenCalledWith(
[mockPathAndCacheKey],
IndexResultType.Delete,
Expand Down
Loading

0 comments on commit 8cdb68c

Please sign in to comment.