test: updating TaskPersistence vitest mocks

Author: edelauna

src/__tests__/nested-delegation-resume.spec.ts

@@ -44,7 +44,8 @@ vi.mock("vscode", () => {
 vi.mock("../core/task-persistence/taskMessages", () => ({
 	readTaskMessages: vi.fn().mockResolvedValue([]),
 }))
-vi.mock("../core/task-persistence", () => ({
+vi.mock("../core/task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../core/task-persistence")>()),
 	readApiMessages: vi.fn().mockResolvedValue([]),
 	saveApiMessages: vi.fn().mockResolvedValue(undefined),
 	saveTaskMessages: vi.fn().mockResolvedValue(undefined),

src/core/task/__tests__/Task.persistence.spec.ts

@@ -73,26 +73,30 @@ vi.mock("p-wait-for", () => ({
 	default: mockPWaitFor,
 }))
 
-vi.mock("../../task-persistence", () => ({
-	saveApiMessages: mockSaveApiMessages,
-	saveTaskMessages: mockSaveTaskMessages,
-	readApiMessages: mockReadApiMessages,
-	readTaskMessages: mockReadTaskMessages,
-	taskMetadata: mockTaskMetadata,
-	TaskHistoryStore: vi.fn().mockImplementation(function () {
-		return {
-			initialize: vi.fn().mockResolvedValue(undefined),
-			dispose: vi.fn(),
-			get: vi.fn(),
-			getAll: vi.fn().mockReturnValue([]),
-			upsert: vi.fn().mockResolvedValue([]),
-			delete: vi.fn().mockResolvedValue(undefined),
-			deleteMany: vi.fn().mockResolvedValue(undefined),
-			reconcile: vi.fn().mockResolvedValue(undefined),
-			initialized: Promise.resolve(),
-		}
-	}),
-}))
+vi.mock("../../task-persistence", async (importOriginal) => {
+	const mod = await importOriginal<typeof import("../../task-persistence")>()
+	return {
+		...mod,
+		saveApiMessages: mockSaveApiMessages,
+		saveTaskMessages: mockSaveTaskMessages,
+		readApiMessages: mockReadApiMessages,
+		readTaskMessages: mockReadTaskMessages,
+		taskMetadata: mockTaskMetadata,
+		TaskHistoryStore: vi.fn().mockImplementation(function () {
+			return {
+				initialize: vi.fn().mockResolvedValue(undefined),
+				dispose: vi.fn(),
+				get: vi.fn(),
+				getAll: vi.fn().mockReturnValue([]),
+				upsert: vi.fn().mockResolvedValue([]),
+				delete: vi.fn().mockResolvedValue(undefined),
+				deleteMany: vi.fn().mockResolvedValue(undefined),
+				reconcile: vi.fn().mockResolvedValue(undefined),
+				initialized: Promise.resolve(),
+			}
+		}),
+	}
+})
 
 vi.mock("vscode", () => {
 	const mockDisposable = { dispose: vi.fn() }

src/core/task/__tests__/Task.throttle.test.ts

@@ -33,7 +33,8 @@ vi.mock("@roo-code/telemetry", () => ({
 }))
 
 // Mock task persistence to avoid disk writes
-vi.mock("../../task-persistence", () => ({
+vi.mock("../../task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../../task-persistence")>()),
 	readApiMessages: vi.fn().mockResolvedValue([]),
 	saveApiMessages: vi.fn().mockResolvedValue(undefined),
 	readTaskMessages: vi.fn().mockResolvedValue([]),

src/core/webview/__tests__/ClineProvider.flicker-free-cancel.spec.ts

@@ -232,24 +232,28 @@ vi.mock("../../../services/skills/SkillsManager", () => ({
 	}),
 }))
 
-vi.mock("../../task-persistence", () => ({
-	TaskHistoryStore: vi.fn().mockImplementation(function () {
-		return {
-			initialize: vi.fn().mockResolvedValue(undefined),
-			dispose: vi.fn(),
-			initialized: Promise.resolve(),
-			get: vi.fn().mockReturnValue(undefined),
-			getAll: vi.fn().mockReturnValue([]),
-			upsert: vi.fn().mockResolvedValue([]),
-			delete: vi.fn().mockResolvedValue(undefined),
-			deleteMany: vi.fn().mockResolvedValue(undefined),
-			migrateFromGlobalState: vi.fn().mockResolvedValue(undefined),
-		}
-	}),
-	readApiMessages: vi.fn().mockResolvedValue([]),
-	saveApiMessages: vi.fn().mockResolvedValue(undefined),
-	saveTaskMessages: vi.fn().mockResolvedValue(undefined),
-}))
+vi.mock("../../task-persistence", async (importOriginal) => {
+	const mod = await importOriginal<typeof import("../../task-persistence")>()
+	return {
+		...mod,
+		TaskHistoryStore: vi.fn().mockImplementation(function () {
+			return {
+				initialize: vi.fn().mockResolvedValue(undefined),
+				dispose: vi.fn(),
+				initialized: Promise.resolve(),
+				get: vi.fn().mockReturnValue(undefined),
+				getAll: vi.fn().mockReturnValue([]),
+				upsert: vi.fn().mockResolvedValue([]),
+				delete: vi.fn().mockResolvedValue(undefined),
+				deleteMany: vi.fn().mockResolvedValue(undefined),
+				migrateFromGlobalState: vi.fn().mockResolvedValue(undefined),
+			}
+		}),
+		readApiMessages: vi.fn().mockResolvedValue([]),
+		saveApiMessages: vi.fn().mockResolvedValue(undefined),
+		saveTaskMessages: vi.fn().mockResolvedValue(undefined),
+	}
+})
 
 describe("ClineProvider flicker-free cancel", () => {
 	let provider: ClineProvider

src/core/webview/__tests__/checkpointRestoreHandler.spec.ts

@@ -5,7 +5,8 @@ import pWaitFor from "p-wait-for"
 import * as vscode from "vscode"
 
 // Mock dependencies
-vi.mock("../../task-persistence", () => ({
+vi.mock("../../task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../../task-persistence")>()),
 	saveTaskMessages: vi.fn(),
 }))
 vi.mock("p-wait-for")

src/core/webview/__tests__/webviewMessageHandler.checkpoint.spec.ts

@@ -6,7 +6,10 @@ import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler"
 import { MessageManager } from "../../message-manager"
 
 // Mock dependencies
-vi.mock("../../task-persistence")
+vi.mock("../../task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../../task-persistence")>()),
+	saveTaskMessages: vi.fn(),
+}))
 vi.mock("../checkpointRestoreHandler")
 vi.mock("p-wait-for", () => ({
 	default: vi.fn(async (condition: () => boolean) => {

src/core/webview/__tests__/webviewMessageHandler.delete.spec.ts

@@ -5,7 +5,8 @@ import { ClineProvider } from "../ClineProvider"
 import { MessageManager } from "../../message-manager"
 
 // Mock the saveTaskMessages function
-vi.mock("../../task-persistence", () => ({
+vi.mock("../../task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../../task-persistence")>()),
 	saveTaskMessages: vi.fn(),
 }))
 

src/core/webview/__tests__/webviewMessageHandler.edit.spec.ts

@@ -22,7 +22,8 @@ vi.mock("vscode", () => ({
 	},
 }))
 
-vi.mock("../../task-persistence", () => ({
+vi.mock("../../task-persistence", async (importOriginal) => ({
+	...(await importOriginal<typeof import("../../task-persistence")>()),
 	saveTaskMessages: vi.fn(),
 }))