Fix ImagenGenerationRequestTests

Author: andrewheard

FirebaseAI/Tests/Unit/Types/Imagen/ImageGenerationParametersTests.swift

@@ -205,6 +205,7 @@ final class ImageGenerationParametersTests: XCTestCase {
     let outputOptions = ImageGenerationOutputOptions(mimeType: mimeType, compressionQuality: nil)
     let addWatermark = false
     let includeRAIReason = true
+    let includeSafetyAttributes = true
     let parameters = ImageGenerationParameters(
       sampleCount: sampleCount,
       storageURI: storageURI,
@@ -215,7 +216,7 @@ final class ImageGenerationParametersTests: XCTestCase {
       outputOptions: outputOptions,
       addWatermark: addWatermark,
       includeResponsibleAIFilterReason: includeRAIReason,
-      includeSafetyAttributes: true
+      includeSafetyAttributes: includeSafetyAttributes
     )
 
     let jsonData = try encoder.encode(parameters)
@@ -226,7 +227,7 @@ final class ImageGenerationParametersTests: XCTestCase {
       "addWatermark" : \(addWatermark),
       "aspectRatio" : "\(aspectRatio)",
       "includeRaiReason" : \(includeRAIReason),
-      "includeSafetyAttributes" : true,
+      "includeSafetyAttributes" : \(includeSafetyAttributes),
       "negativePrompt" : "\(negativePrompt)",
       "outputOptions" : {
         "mimeType" : "\(mimeType)"
@@ -253,7 +254,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: addWatermark,
-      includeResponsibleAIFilterReason: nil
+      includeResponsibleAIFilterReason: nil,
+      includeSafetyAttributes: nil
     )
 
     let jsonData = try encoder.encode(parameters)
@@ -279,7 +281,8 @@ final class ImageGenerationParametersTests: XCTestCase {
       personGeneration: nil,
       outputOptions: nil,
       addWatermark: nil,
-      includeResponsibleAIFilterReason: nil
+      includeResponsibleAIFilterReason: nil,
+      includeSafetyAttributes: nil
     )
 
     let jsonData = try encoder.encode(parameters)

FirebaseAI/Tests/Unit/Types/Imagen/ImagenGenerationRequestTests.swift

@@ -25,6 +25,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
   let aspectRatio = "16:9"
   let safetyFilterLevel = "block_low_and_above"
   let includeResponsibleAIFilterReason = true
+  let includeSafetyAttributes = true
   lazy var parameters = ImageGenerationParameters(
     sampleCount: sampleCount,
     storageURI: nil,
@@ -34,7 +35,8 @@ final class ImagenGenerationRequestTests: XCTestCase {
     personGeneration: nil,
     outputOptions: nil,
     addWatermark: nil,
-    includeResponsibleAIFilterReason: includeResponsibleAIFilterReason
+    includeResponsibleAIFilterReason: includeResponsibleAIFilterReason,
+    includeSafetyAttributes: includeSafetyAttributes
   )
   let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
 
@@ -108,6 +110,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
       "parameters" : {
         "aspectRatio" : "\(aspectRatio)",
         "includeRaiReason" : \(includeResponsibleAIFilterReason),
+        "includeSafetyAttributes" : \(includeSafetyAttributes),
         "safetySetting" : "\(safetyFilterLevel)",
         "sampleCount" : \(sampleCount)
       }
@@ -137,6 +140,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
       "parameters" : {
         "aspectRatio" : "\(aspectRatio)",
         "includeRaiReason" : \(includeResponsibleAIFilterReason),
+        "includeSafetyAttributes" : \(includeSafetyAttributes),
         "safetySetting" : "\(safetyFilterLevel)",
         "sampleCount" : \(sampleCount)
       }