Refactor: Use config.getGeminiClient() for GeminiClient instantiation (#715)

1 files changed, 6 insertions, 0 deletions

diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts
index c94edfd1..4646f30a 100644
--- a/packages/core/src/tools/write-file.test.ts
+++ b/packages/core/src/tools/write-file.test.ts
@@ -53,6 +53,7 @@ const mockConfigInternal = {
   getTargetDir: () => rootDir,
   getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
   setApprovalMode: vi.fn(),
+  getGeminiClient: vi.fn(), // Initialize as a plain mock function
   getApiKey: () => 'test-key',
   getModel: () => 'test-model',
   getSandbox: () => false,
@@ -97,6 +98,11 @@ describe('WriteFileTool', () => {
     ) as Mocked<GeminiClient>;
     vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance);
 
+    // Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient
+    mockConfigInternal.getGeminiClient.mockReturnValue(
+      mockGeminiClientInstance,
+    );
+
     tool = new WriteFileTool(mockConfig);
 
     // Reset mocks before each test