From d3c70e63da6f07e30c5ebfdfb2e59fda6b29ad7d Mon Sep 17 00:00:00 2001
From: Rodrigo Lazo Paz <rlazo@google.com>
Date: Mon, 28 Apr 2025 12:59:37 -0400
Subject: [PATCH] [ai] Address most tests failures

Citation test still times out, but this address all other breakages.
---
 .../google/firebase/ai/DevAPIUnarySnapshotTests.kt  | 13 +------------
 .../test/java/com/google/firebase/ai/util/tests.kt  |  2 +-
 firebase-ai/update_responses.sh                     |  2 +-
 3 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIUnarySnapshotTests.kt b/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIUnarySnapshotTests.kt
index 7eb06a702fb..0517d350c2d 100644
--- a/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIUnarySnapshotTests.kt
+++ b/firebase-ai/src/test/java/com/google/firebase/ai/DevAPIUnarySnapshotTests.kt
@@ -57,24 +57,13 @@ internal class DevAPIUnarySnapshotTests {
       }
     }
 
-  @Test
-  fun `prompt blocked for safety`() =
-    goldenDevAPIUnaryFile("unary-failure-prompt-blocked-safety.txt") {
-      withTimeout(testTimeout) {
-        shouldThrow<ResponseStoppedException> { model.generateContent("prompt") } should
-          {
-            it.response.candidates[0].finishReason shouldBe FinishReason.MAX_TOKENS
-          }
-      }
-    }
-
   @Test
   fun `response blocked for safety`() =
     goldenDevAPIUnaryFile("unary-failure-finish-reason-safety.txt") {
       withTimeout(testTimeout) {
         shouldThrow<ResponseStoppedException> { model.generateContent("prompt") } should
           {
-            it.response.candidates[0].finishReason shouldBe FinishReason.MAX_TOKENS
+            it.response.candidates[0].finishReason shouldBe FinishReason.SAFETY
           }
       }
     }
diff --git a/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt b/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt
index 44b7a1dff21..4c325faa884 100644
--- a/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt
+++ b/firebase-ai/src/test/java/com/google/firebase/ai/util/tests.kt
@@ -248,7 +248,7 @@ internal fun goldenDevAPIUnaryFile(
   name: String,
   httpStatusCode: HttpStatusCode = HttpStatusCode.OK,
   block: CommonTest,
-) = goldenUnaryFile("developerapi/$name", httpStatusCode, GenerativeBackend.googleAI(), block)
+) = goldenUnaryFile("googleai/$name", httpStatusCode, GenerativeBackend.googleAI(), block)
 
 /**
  * Loads a *Golden File* from the resource directory.
diff --git a/firebase-ai/update_responses.sh b/firebase-ai/update_responses.sh
index 28b0be1a25a..e814fa9eedd 100755
--- a/firebase-ai/update_responses.sh
+++ b/firebase-ai/update_responses.sh
@@ -17,7 +17,7 @@
 # This script replaces mock response files for Vertex AI unit tests with a fresh
 # clone of the shared repository of Vertex AI test data.
 
-RESPONSES_VERSION='v7.*' # The major version of mock responses to use
+RESPONSES_VERSION='v10.*' # The major version of mock responses to use
 REPO_NAME="vertexai-sdk-test-data"
 REPO_LINK="https://github.com/FirebaseExtended/$REPO_NAME.git"