From 8ab74ef1bb4bd1c475596088d9d3b52e0a9c5ca7 Mon Sep 17 00:00:00 2001
From: "N. Taylor Mullen" <ntaylormullen@google.com>
Date: Mon, 2 Jun 2025 22:30:52 -0700
Subject: Refactor: Use config.getGeminiClient() for GeminiClient instantiation
 (#715)

---
 packages/cli/src/ui/hooks/useGeminiStream.test.tsx | 27 ++++++++++++++--------
 1 file changed, 18 insertions(+), 9 deletions(-)

(limited to 'packages/cli/src/ui/hooks/useGeminiStream.test.tsx')

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index d46fab9e..4dfea5e0 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -25,20 +25,20 @@ const mockSendMessageStream = vi
   .mockReturnValue((async function* () {})());
 const mockStartChat = vi.fn();
 
-vi.mock('@gemini-code/core', async (importOriginal) => {
-  const actualCoreModule = (await importOriginal()) as any;
-  const MockedGeminiClientClass = vi.fn().mockImplementation(function (
-    this: any,
-    _config: any,
-  ) {
+const MockedGeminiClientClass = vi.hoisted(() =>
+  vi.fn().mockImplementation(function (this: any, _config: any) {
     // _config
     this.startChat = mockStartChat;
     this.sendMessageStream = mockSendMessageStream;
-  });
+  }),
+);
+
+vi.mock('@gemini-code/core', async (importOriginal) => {
+  const actualCoreModule = (await importOriginal()) as any;
   return {
     ...(actualCoreModule || {}),
-    GeminiClient: MockedGeminiClientClass,
-    // GeminiChat will be from actualCoreModule if it exists, otherwise undefined
+    GeminiClient: MockedGeminiClientClass, // Export the class for type checking or other direct uses
+    Config: actualCoreModule.Config, // Ensure Config is passed through
   };
 });
 
@@ -235,6 +235,14 @@ describe('useGeminiStream', () => {
 
     mockAddItem = vi.fn();
     mockSetShowHelp = vi.fn();
+    // Define the mock for getGeminiClient
+    const mockGetGeminiClient = vi.fn().mockImplementation(() => {
+      // MockedGeminiClientClass is defined in the module scope by the previous change.
+      // It will use the mockStartChat and mockSendMessageStream that are managed within beforeEach.
+      const clientInstance = new MockedGeminiClientClass(mockConfig);
+      return clientInstance;
+    });
+
     mockConfig = {
       apiKey: 'test-api-key',
       model: 'gemini-pro',
@@ -258,6 +266,7 @@ describe('useGeminiStream', () => {
       getToolRegistry: vi.fn(
         () => ({ getToolSchemaList: vi.fn(() => []) }) as any,
       ),
+      getGeminiClient: mockGetGeminiClient,
     } as unknown as Config;
     mockOnDebugMessage = vi.fn();
     mockHandleSlashCommand = vi.fn().mockReturnValue(false);
-- 
cgit v1.2.3