From 48e287892fc3db00add5e00b95d9d2631234664d Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Thu, 21 Mar 2024 13:04:25 +0700
Subject: [PATCH] test: use unique tmp dir for storage tests and wait to clean
 VectorStoreIndex files

---
 packages/core/tests/StorageContext.test.ts    | 11 ++++++----
 .../core/tests/indices/SummaryIndex.test.ts   |  9 ++++++--
 .../tests/indices/VectorStoreIndex.test.ts    | 21 +++++++++++++------
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/packages/core/tests/StorageContext.test.ts b/packages/core/tests/StorageContext.test.ts
index bfbb4924f..95dfc7c1e 100644
--- a/packages/core/tests/StorageContext.test.ts
+++ b/packages/core/tests/StorageContext.test.ts
@@ -3,6 +3,9 @@ import {
   type StorageContext,
 } from "llamaindex/storage/StorageContext";
 import { existsSync, rmSync } from "node:fs";
+import { mkdtemp } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 import {
   afterAll,
   beforeAll,
@@ -12,7 +15,7 @@ import {
   vi,
   vitest,
 } from "vitest";
-
+const testDir = await mkdtemp(join(tmpdir(), "test-"));
 vitest.spyOn(console, "error");
 
 describe("StorageContext", () => {
@@ -20,18 +23,18 @@ describe("StorageContext", () => {
 
   beforeAll(async () => {
     storageContext = await storageContextFromDefaults({
-      persistDir: "/tmp/test_dir",
+      persistDir: testDir,
     });
   });
 
   test("initializes", async () => {
     vi.mocked(console.error).mockImplementation(() => {}); // silence console.error
 
-    expect(existsSync("/tmp/test_dir")).toBe(true);
+    expect(existsSync(testDir)).toBe(true);
     expect(storageContext).toBeDefined();
   });
 
   afterAll(() => {
-    rmSync("/tmp/test_dir", { recursive: true });
+    rmSync(testDir, { recursive: true });
   });
 });
diff --git a/packages/core/tests/indices/SummaryIndex.test.ts b/packages/core/tests/indices/SummaryIndex.test.ts
index 89277f6e8..f43df0ce7 100644
--- a/packages/core/tests/indices/SummaryIndex.test.ts
+++ b/packages/core/tests/indices/SummaryIndex.test.ts
@@ -7,8 +7,13 @@ import {
   type StorageContext,
 } from "llamaindex";
 import { rmSync } from "node:fs";
+import { mkdtemp } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 
+const testDir = await mkdtemp(join(tmpdir(), "test-"));
+
 vi.mock("llamaindex/llm/open_ai", () => {
   return {
     getOpenAISession: vi.fn().mockImplementation(() => null),
@@ -24,7 +29,7 @@ describe("SummaryIndex", () => {
   beforeAll(async () => {
     serviceContext = mockServiceContext();
     storageContext = await storageContextFromDefaults({
-      persistDir: "/tmp/test_dir",
+      persistDir: testDir,
     });
   });
 
@@ -43,6 +48,6 @@ describe("SummaryIndex", () => {
   });
 
   afterAll(() => {
-    rmSync("/tmp/test_dir", { recursive: true });
+    rmSync(testDir, { recursive: true });
   });
 });
diff --git a/packages/core/tests/indices/VectorStoreIndex.test.ts b/packages/core/tests/indices/VectorStoreIndex.test.ts
index b79af7311..43b6aa023 100644
--- a/packages/core/tests/indices/VectorStoreIndex.test.ts
+++ b/packages/core/tests/indices/VectorStoreIndex.test.ts
@@ -4,8 +4,13 @@ import {
   VectorStoreIndex,
   storageContextFromDefaults,
 } from "llamaindex";
-import { beforeAll, describe, expect, test, vi } from "vitest";
-import { mockServiceContext } from "../utility/mockServiceContext.js";
+import { rmSync } from "node:fs";
+import { mkdtemp } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterAll, beforeAll, describe, expect, test, vi } from "vitest";
+
+const testDir = await mkdtemp(join(tmpdir(), "test-"));
 
 vi.mock("llamaindex/llm/open_ai", () => {
   return {
@@ -13,6 +18,8 @@ vi.mock("llamaindex/llm/open_ai", () => {
   };
 });
 
+import { mockServiceContext } from "../utility/mockServiceContext.js";
+
 describe.sequential("VectorStoreIndex", () => {
   let serviceContext: ServiceContext;
   let storageContext: StorageContext;
@@ -24,7 +31,7 @@ describe.sequential("VectorStoreIndex", () => {
   beforeAll(async () => {
     serviceContext = mockServiceContext();
     storageContext = await storageContextFromDefaults({
-      persistDir: "/tmp/test_dir",
+      persistDir: testDir,
     });
     testStrategy = async (
       // strategy?: DocStoreStrategy,
@@ -55,7 +62,9 @@ describe.sequential("VectorStoreIndex", () => {
   //   expect(entries[0]).toBe(entries[1]);
   // });
 
-  // afterAll(() => {
-  //   rmSync("/tmp/test_dir", { recursive: true });
-  // });
+  afterAll(async () => {
+    // TODO: VectorStoreIndex.fromDocuments running twice is causing a cleanup issue
+    await new Promise((resolve) => setTimeout(resolve, 100));
+    rmSync(testDir, { recursive: true });
+  });
 });
-- 
GitLab