1 files changed, 6 insertions, 1 deletions
diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
index 922d20fb..d8d9c10a 100644
--- a/packages/core/src/code_assist/server.test.ts
+++ b/packages/core/src/code_assist/server.test.ts
@@ -133,11 +133,16 @@ describe('CodeAssistServer', () => {
   it('should return 0 for countTokens', async () => {
     const auth = new OAuth2Client();
     const server = new CodeAssistServer(auth, 'test-project');
+    const mockResponse = {
+      totalTokens: 100,
+    };
+    vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse);
+
     const response = await server.countTokens({
       model: 'test-model',
       contents: [{ role: 'user', parts: [{ text: 'request' }] }],
     });
-    expect(response.totalTokens).toBe(0);
+    expect(response.totalTokens).toBe(100);
   });
 
   it('should throw an error for embedContent', async () => {