From 329f413870b59a61d6d7b3df8a9eaf5182ea2cfc Mon Sep 17 00:00:00 2001
From: can1357 <me@can.ac>
Date: Thu, 26 Mar 2026 14:51:30 +0100
Subject: [PATCH] test(ai): refactored auth storage test to use spyOn for
 cleaner mocks

- Refactored auth storage test to use vi.spyOn() instead of vi.mock() for cleaner mock management.
- Simplified mock type definitions by leveraging bun:test's Mock type import.
---
 .../test/auth-storage-api-key-login.test.ts   | 20 +++++--------------
 1 file changed, 5 insertions(+), 15 deletions(-)

diff --git a/packages/ai/test/auth-storage-api-key-login.test.ts b/packages/ai/test/auth-storage-api-key-login.test.ts
index b4f52c4ae..09a1373ec 100644
--- a/packages/ai/test/auth-storage-api-key-login.test.ts
+++ b/packages/ai/test/auth-storage-api-key-login.test.ts
@@ -1,22 +1,11 @@
 import { Database } from "bun:sqlite";
-import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
+import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from "bun:test";
 import * as fs from "node:fs/promises";
 import * as os from "node:os";
 import * as path from "node:path";
 
-vi.mock("../src/utils/oauth/kagi", () => ({
-	loginKagi: vi.fn(),
-}));
-
 import { AuthCredentialStore, AuthStorage } from "../src/auth-storage";
-import { loginKagi } from "../src/utils/oauth/kagi";
-
-type MockedApiKeyLogin = {
-	mockReset(): void;
-	mockResolvedValueOnce(value: string): MockedApiKeyLogin;
-};
-
-const mockedLoginKagi = loginKagi as typeof loginKagi & MockedApiKeyLogin;
+import * as kagiModule from "../src/utils/oauth/kagi";
 
 function countCredentialRows(dbPath: string, provider: string): number {
 	const db = new Database(dbPath, { readonly: true });
@@ -35,13 +24,14 @@ describe("AuthStorage api-key login replacement", () => {
 	let dbPath = "";
 	let store: AuthCredentialStore | null = null;
 	let authStorage: AuthStorage | null = null;
+	let loginKagiSpy: Mock<typeof kagiModule.loginKagi>;
 
 	beforeEach(async () => {
 		tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-api-key-login-"));
 		dbPath = path.join(tempDir, "agent.db");
 		store = await AuthCredentialStore.open(dbPath);
 		authStorage = new AuthStorage(store);
-		mockedLoginKagi.mockReset();
+		loginKagiSpy = vi.spyOn(kagiModule, "loginKagi");
 	});
 
 	afterEach(async () => {
@@ -59,7 +49,7 @@ describe("AuthStorage api-key login replacement", () => {
 	it("reuses the stored api-key row when re-login returns the same key", async () => {
 		if (!store || !authStorage || !dbPath) throw new Error("test setup failed");
 
-		mockedLoginKagi.mockResolvedValueOnce("same-kagi-key").mockResolvedValueOnce("same-kagi-key");
+		loginKagiSpy.mockResolvedValueOnce("same-kagi-key").mockResolvedValueOnce("same-kagi-key");
 
 		const controller = {
 			onAuth: () => {},