From ef72ac33de7cfc9586605f8c8779952ad923a50d Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 13 Feb 2025 11:58:15 +0000
Subject: [PATCH] fix: mock both language models and safe generator

Co-Authored-By: Han Xiao <han.xiao@jina.ai>
---
 src/tools/__tests__/evaluator.test.ts | 32 ++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/tools/__tests__/evaluator.test.ts b/src/tools/__tests__/evaluator.test.ts
index 43628355..035bda28 100644
--- a/src/tools/__tests__/evaluator.test.ts
+++ b/src/tools/__tests__/evaluator.test.ts
@@ -66,9 +66,9 @@ jest.mock('@ai-sdk/google', () => ({
   GoogleChatLanguageModel: jest.fn()
 }));
 
-// Mock the safe generator to prevent actual API calls
-jest.mock('../../utils/safe-generator', () => ({
-  ObjectGeneratorSafe: jest.fn().mockImplementation(() => ({
+// Mock the language models
+jest.mock('@ai-sdk/google', () => ({
+  GoogleChatLanguageModel: jest.fn().mockImplementation(() => ({
     generateObject: jest.fn().mockImplementation(async () => ({
       object: mockEvalResponse,
       usage: { totalTokens: 0, promptTokens: 0, completionTokens: 0 }
@@ -76,6 +76,32 @@ jest.mock('../../utils/safe-generator', () => ({
   }))
 }));
 
+// Mock the safe generator to prevent actual API calls
+jest.mock('../../utils/safe-generator', () => ({
+  ObjectGeneratorSafe: jest.fn().mockImplementation(() => ({
+    generateObject: jest.fn().mockImplementation(async (params: any) => {
+      const { schema } = params;
+      // Return different responses based on the schema
+      if (schema.description?.includes('question requires freshness')) {
+        return {
+          object: {
+            needsFreshness: false,
+            needsPlurality: false,
+            reasoning: 'Simple test question',
+            languageStyle: 'plain English'
+          },
+          usage: { totalTokens: 0, promptTokens: 0, completionTokens: 0 }
+        };
+      }
+      // Default response for other schemas
+      return {
+        object: mockEvalResponse,
+        usage: { totalTokens: 0, promptTokens: 0, completionTokens: 0 }
+      };
+    })
+  }))
+}));
+
 // Mock safe generator to prevent actual API calls
 jest.mock('../../utils/safe-generator', () => ({
   ObjectGeneratorSafe: jest.fn().mockImplementation(() => ({