From c81946930e840da14039b690d2135e6c65ab24e4 Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Fri, 15 Mar 2024 15:20:57 +0700
Subject: [PATCH] test: fix openai mock
---
packages/core/tests/indices/SummaryIndex.test.ts | 9 ++++++++-
packages/core/tests/indices/VectorStoreIndex.test.ts | 8 +++++++-
packages/core/tests/objects/ObjectIndex.test.ts | 9 ++++++++-
packages/core/tests/utility/mockServiceContext.ts | 9 ---------
4 files changed, 23 insertions(+), 12 deletions(-)
diff --git a/packages/core/tests/indices/SummaryIndex.test.ts b/packages/core/tests/indices/SummaryIndex.test.ts
index 846a61e9d..89277f6e8 100644
--- a/packages/core/tests/indices/SummaryIndex.test.ts
+++ b/packages/core/tests/indices/SummaryIndex.test.ts
@@ -7,7 +7,14 @@ import {
type StorageContext,
} from "llamaindex";
import { rmSync } from "node:fs";
-import { afterAll, beforeAll, describe, expect, it } from "vitest";
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
+
+vi.mock("llamaindex/llm/open_ai", () => {
+ return {
+ getOpenAISession: vi.fn().mockImplementation(() => null),
+ };
+});
+
import { mockServiceContext } from "../utility/mockServiceContext.js";
describe("SummaryIndex", () => {
diff --git a/packages/core/tests/indices/VectorStoreIndex.test.ts b/packages/core/tests/indices/VectorStoreIndex.test.ts
index 4e9f3694a..b79af7311 100644
--- a/packages/core/tests/indices/VectorStoreIndex.test.ts
+++ b/packages/core/tests/indices/VectorStoreIndex.test.ts
@@ -4,9 +4,15 @@ import {
VectorStoreIndex,
storageContextFromDefaults,
} from "llamaindex";
-import { beforeAll, describe, expect, test } from "vitest";
+import { beforeAll, describe, expect, test, vi } from "vitest";
import { mockServiceContext } from "../utility/mockServiceContext.js";
+vi.mock("llamaindex/llm/open_ai", () => {
+ return {
+ getOpenAISession: vi.fn().mockImplementation(() => null),
+ };
+});
+
describe.sequential("VectorStoreIndex", () => {
let serviceContext: ServiceContext;
let storageContext: StorageContext;
diff --git a/packages/core/tests/objects/ObjectIndex.test.ts b/packages/core/tests/objects/ObjectIndex.test.ts
index 08203f50f..356657919 100644
--- a/packages/core/tests/objects/ObjectIndex.test.ts
+++ b/packages/core/tests/objects/ObjectIndex.test.ts
@@ -5,7 +5,14 @@ import {
SimpleToolNodeMapping,
VectorStoreIndex,
} from "llamaindex";
-import { beforeAll, describe, expect, test } from "vitest";
+import { beforeAll, describe, expect, test, vi } from "vitest";
+
+vi.mock("llamaindex/llm/open_ai", () => {
+ return {
+ getOpenAISession: vi.fn().mockImplementation(() => null),
+ };
+});
+
import { mockServiceContext } from "../utility/mockServiceContext.js";
describe("ObjectIndex", () => {
diff --git a/packages/core/tests/utility/mockServiceContext.ts b/packages/core/tests/utility/mockServiceContext.ts
index 03b0893b8..f6843069c 100644
--- a/packages/core/tests/utility/mockServiceContext.ts
+++ b/packages/core/tests/utility/mockServiceContext.ts
@@ -9,15 +9,6 @@ import {
mockLlmGeneration,
} from "../utility/mockOpenAI.js";
-import { vi } from "vitest";
-
-// Mock the OpenAI getOpenAISession function during testing
-vi.mock("llamaindex/llm/open_ai", () => {
- return {
- getOpenAISession: vi.fn().mockImplementation(() => null),
- };
-});
-
export function mockServiceContext() {
const embeddingModel = new OpenAIEmbedding();
const llm = new OpenAI();
--
GitLab