From 3d2f294f6a5008063ad3fad04a45c173f5aba542 Mon Sep 17 00:00:00 2001
From: Lex Christopherson <lex@glittercowboy.com>
Date: Mon, 16 Mar 2026 11:56:46 -0600
Subject: [PATCH] fix: google-search OAuth test mock and Windows path separator
 in smoke test

- google-search test: mock getApiKeyForProvider to return JSON string
  matching real OAuth provider behavior (token+projectId), instead of
  using AuthStorage.inMemory which bypasses the OAuth getApiKey transform
- smoke test: split on /[/\\]/ for Windows path separator compatibility

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/tests/extension-smoke.test.ts          |  2 +-
 src/tests/google-search-auth.repro.test.ts | 54 ++++++++++++----------
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/src/tests/extension-smoke.test.ts b/src/tests/extension-smoke.test.ts
index f8b9f4f37..795bb941f 100644
--- a/src/tests/extension-smoke.test.ts
+++ b/src/tests/extension-smoke.test.ts
@@ -39,7 +39,7 @@ test("all bundled extensions can be imported without throwing", async () => {
 
   for (const entryPath of entryPaths) {
     const relPath = entryPath.slice(extensionsDir.length + 1);
-    const extName = relPath.split("/")[0].replace(/\.ts$/, "");
+    const extName = relPath.split(/[/\\]/)[0].replace(/\.ts$/, "");
 
     if (SKIP_EXTENSIONS.has(extName)) {
       skipped++;
diff --git a/src/tests/google-search-auth.repro.test.ts b/src/tests/google-search-auth.repro.test.ts
index 2b3bab22a..69198845b 100644
--- a/src/tests/google-search-auth.repro.test.ts
+++ b/src/tests/google-search-auth.repro.test.ts
@@ -1,16 +1,13 @@
 import test from "node:test";
 import assert from "node:assert/strict";
-import { AuthStorage, ModelRegistry } from "../../packages/pi-coding-agent/src/index.js";
 import googleSearchExtension from "../resources/extensions/google-search/index.ts";
 
 function createMockPI() {
   const handlers: any[] = [];
-  const notifications: any[] = [];
   let registeredTool: any = null;
 
   return {
     handlers,
-    notifications,
     registeredTool,
     on(event: string, handler: any) {
       handlers.push({ event, handler });
@@ -28,11 +25,23 @@ function createMockPI() {
   };
 }
 
+/**
+ * Build a mock modelRegistry whose getApiKeyForProvider returns the given
+ * JSON string (matching what the real OAuth provider's getApiKey produces).
+ */
+function mockModelRegistry(oauthJson?: string) {
+  return {
+    authStorage: {
+      hasAuth: async (_id: string) => !!oauthJson,
+    },
+    getApiKeyForProvider: async (_provider: string) => oauthJson,
+  };
+}
+
 test("fix: google-search uses OAuth if GEMINI_API_KEY is missing", async () => {
   const originalKey = process.env.GEMINI_API_KEY;
   delete process.env.GEMINI_API_KEY;
 
-  // Mock fetch
   const originalFetch = global.fetch;
   (global as any).fetch = async (url: string, options: any) => {
     assert.ok(url.includes("cloudcode-pa.googleapis.com"), "Should use Cloud Code Assist endpoint");
@@ -43,23 +52,29 @@ test("fix: google-search uses OAuth if GEMINI_API_KEY is missing", async () => {
         response: {
           candidates: [{ content: { parts: [{ text: "Mocked AI Answer" }] } }]
         }
-      })
+      }),
+      text: async () => JSON.stringify({
+        response: {
+          candidates: [{ content: { parts: [{ text: "Mocked AI Answer" }] } }]
+        }
+      }),
     };
   };
 
   try {
     const pi = createMockPI();
     googleSearchExtension(pi as any);
-    const authStorage = AuthStorage.inMemory({
-      "google-gemini-cli": { type: "oauth", access: "mock-token", projectId: "mock-project" }
-    });
-    const modelRegistry = new ModelRegistry(authStorage);
-    const mockCtx = { ui: { notify() {} }, modelRegistry };
+
+    const oauthJson = JSON.stringify({ token: "mock-token", projectId: "mock-project" });
+    const mockCtx = {
+      ui: { notify() {} },
+      modelRegistry: mockModelRegistry(oauthJson),
+    };
 
     await pi.fire("session_start", {}, mockCtx);
     const registeredTool = (pi as any).registeredTool;
     const result = await registeredTool.execute("call-1", { query: "test" }, new AbortController().signal, () => {}, mockCtx);
-    
+
     assert.equal(result.isError, undefined);
     assert.ok(result.content[0].text.includes("Mocked AI Answer"));
   } finally {
@@ -75,12 +90,11 @@ test("google-search warns if NO authentication is present", async () => {
   try {
     const pi = createMockPI();
     googleSearchExtension(pi as any);
-    const authStorage = AuthStorage.inMemory({}); // No OAuth
-    const modelRegistry = new ModelRegistry(authStorage);
+
     const notifications: any[] = [];
     const mockCtx = {
       ui: { notify(msg: string, level: string) { notifications.push({ msg, level }); } },
-      modelRegistry
+      modelRegistry: mockModelRegistry(undefined),
     };
 
     await pi.fire("session_start", {}, mockCtx);
@@ -102,23 +116,15 @@ test("google-search uses GEMINI_API_KEY if present (precedence)", async () => {
   try {
     const pi = createMockPI();
     googleSearchExtension(pi as any);
-    
-    // Even if OAuth is available, it should prefer the API Key
-    const authStorage = AuthStorage.inMemory({
-      "google-gemini-cli": { type: "oauth", access: "should-not-be-used", projectId: "mock-project" }
-    });
-    const modelRegistry = new ModelRegistry(authStorage);
+
     const notifications: any[] = [];
     const mockCtx = {
       ui: { notify(msg: string, level: string) { notifications.push({ msg, level }); } },
-      modelRegistry
+      modelRegistry: mockModelRegistry(JSON.stringify({ token: "should-not-be-used", projectId: "mock-project" })),
     };
 
     await pi.fire("session_start", {}, mockCtx);
     assert.equal(notifications.length, 0, "Should NOT notify if API Key is present");
-
-    // We don't easily mock the @google/genai client here without more effort, 
-    // but we've verified the logic branches.
   } finally {
     delete process.env.GEMINI_API_KEY;
   }