[AI] Update integration tests to Gemini 3.1 (#15898)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: paulb777

FirebaseAI/Tests/TestApp/Sources/Constants.swift

@@ -29,6 +29,6 @@ public enum ModelNames {
   public static let gemini2_5_FlashLive = "gemini-live-2.5-flash-native-audio"
   public static let gemini2_5_FlashLivePreview = "gemini-2.5-flash-native-audio-preview-12-2025"
   public static let gemini2_5_Pro = "gemini-2.5-pro"
-  public static let gemini3FlashPreview = "gemini-3-flash-preview"
+  public static let gemini3_1_FlashLitePreview = "gemini-3.1-flash-lite-preview"
   public static let gemma3_4B = "gemma-3-4b-it"
 }

FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

@@ -51,8 +51,8 @@ struct GenerateContentIntegrationTests {
     (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2FlashLite),
     (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2FlashLite),
     (InstanceConfig.vertexAI_v1beta_global_appCheckLimitedUse, ModelNames.gemini2FlashLite),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini3FlashPreview),
-    (InstanceConfig.googleAI_v1beta_appCheckLimitedUse, ModelNames.gemini3FlashPreview),
+    (InstanceConfig.googleAI_v1beta, ModelNames.gemini3_1_FlashLitePreview),
+    (InstanceConfig.googleAI_v1beta_appCheckLimitedUse, ModelNames.gemini3_1_FlashLitePreview),
     (InstanceConfig.googleAI_v1beta, ModelNames.gemma3_4B),
     (InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemma3_4B),
     // Note: The following configs are commented out for easy one-off manual testing.
@@ -82,12 +82,8 @@ struct GenerateContentIntegrationTests {
     let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
     #expect(promptTokensDetails.modality == .text)
     #expect(promptTokensDetails.tokenCount == usageMetadata.promptTokenCount)
-    if modelName.hasPrefix("gemini-3") {
-      // For gemini-3 models, the thoughtsTokenCount can vary slightly between runs.
-      #expect(usageMetadata.thoughtsTokenCount >= 64)
-    } else {
-      #expect(usageMetadata.thoughtsTokenCount == 0)
-    }
+    // No thoughts in Flash Lite.
+    #expect(usageMetadata.thoughtsTokenCount == 0)
     // The fields `candidatesTokenCount` and `candidatesTokensDetails` are not included when using
     // Gemma models.
     if modelName.hasPrefix("gemini-3") {
@@ -176,13 +172,37 @@ struct GenerateContentIntegrationTests {
       (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
         thinkingBudget: 32768, includeThoughts: true
       )),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingLevel: .minimal)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingLevel: .low)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingLevel: .medium)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingLevel: .high)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingBudget: 128)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingBudget: 32768)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingLevel: .minimal)
+      ),
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingLevel: .low)
+      ),
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingLevel: .medium)
+      ),
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingLevel: .high)
+      ),
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingBudget: 128)
+      ),
+      (
+        .googleAI_v1beta,
+        ModelNames.gemini3_1_FlashLitePreview,
+        ThinkingConfig(thinkingBudget: 32768)
+      ),
+      (.googleAI_v1beta, ModelNames.gemini3_1_FlashLitePreview, ThinkingConfig(
         thinkingBudget: 32768, includeThoughts: true
       )),
       // Note: The following configs are commented out for easy one-off manual testing.
@@ -295,8 +315,8 @@ struct GenerateContentIntegrationTests {
       (.googleAI_v1beta, ModelNames.gemini2_5_Pro, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
       )),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(thinkingBudget: -1)),
-      (.googleAI_v1beta, ModelNames.gemini3FlashPreview, ThinkingConfig(
+      (.googleAI_v1beta, ModelNames.gemini3_1_FlashLitePreview, ThinkingConfig(thinkingBudget: -1)),
+      (.googleAI_v1beta, ModelNames.gemini3_1_FlashLitePreview, ThinkingConfig(
         thinkingBudget: -1, includeThoughts: true
       )),
     ] as [(InstanceConfig, String, ThinkingConfig)]
@@ -502,8 +522,11 @@ struct GenerateContentIntegrationTests {
 
   @Test(arguments: [
     (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2FlashLite),
-    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini3FlashPreview),
-    (InstanceConfig.vertexAI_v1beta_global_appCheckLimitedUse, ModelNames.gemini3FlashPreview),
+    (InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini3_1_FlashLitePreview),
+    (
+      InstanceConfig.vertexAI_v1beta_global_appCheckLimitedUse,
+      ModelNames.gemini3_1_FlashLitePreview
+    ),
     (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_FlashLite),
     (InstanceConfig.googleAI_v1beta_appCheckLimitedUse, ModelNames.gemini2_5_FlashLite),
     (InstanceConfig.googleAI_v1beta, ModelNames.gemma3_4B),
@@ -556,7 +579,11 @@ struct GenerateContentIntegrationTests {
     #expect(promptTextPart.text == prompt)
     let modelHistory = try #require(chat.history.last)
     #expect(modelHistory.role == "model")
-    #expect(modelHistory.parts.count == 1)
+    if modelName.hasPrefix("gemini-3.1-") {
+      #expect(modelHistory.parts.count == 2)
+    } else {
+      #expect(modelHistory.parts.count == 1)
+    }
     let modelTextPart = try #require(modelHistory.parts.first as? TextPart)
     let modelJSONData = try #require(modelTextPart.text.data(using: .utf8))
     let response = try JSONDecoder().decode([String].self, from: modelJSONData)

FirebaseAI/Tests/TestApp/Tests/Integration/ImplicitCacheTests.swift

@@ -30,7 +30,7 @@ struct ImplicitCacheTests {
     (InstanceConfig.vertexAI_v1beta, ModelNames.gemini2_5_Flash),
     (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Flash),
     (InstanceConfig.googleAI_v1beta, ModelNames.gemini2_5_Pro),
-    (InstanceConfig.googleAI_v1beta, ModelNames.gemini3FlashPreview),
+    (InstanceConfig.googleAI_v1beta, ModelNames.gemini3_1_FlashLitePreview),
   ])
   func implicitCaching(_ config: InstanceConfig, modelName: String) async throws {
     let model = FirebaseAI.componentInstance(config).generativeModel(