[Vertex AI] Add error handling for decoding ImagenInlineImage from b64 (#14458)

Author: andrewheard

FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineImage.swift

@@ -26,12 +26,8 @@ public struct ImagenInlineImage {
   /// The image data in PNG or JPEG format.
   public let data: Data
 
-  init(mimeType: String, bytesBase64Encoded: String) {
+  init(mimeType: String, data: Data) {
     self.mimeType = mimeType
-    guard let data = Data(base64Encoded: bytesBase64Encoded) else {
-      // TODO(#14221): Add error handling for invalid base64 bytes.
-      fatalError("Creating a `Data` from `bytesBase64Encoded` failed.")
-    }
     self.data = data
   }
 }
@@ -65,6 +61,13 @@ extension ImagenInlineImage: Decodable {
     let container = try decoder.container(keyedBy: CodingKeys.self)
     let mimeType = try container.decode(String.self, forKey: .mimeType)
     let bytesBase64Encoded = try container.decode(String.self, forKey: .bytesBase64Encoded)
-    self.init(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded)
+    guard let data = Data(base64Encoded: bytesBase64Encoded) else {
+      throw DecodingError.dataCorruptedError(
+        forKey: .bytesBase64Encoded,
+        in: container,
+        debugDescription: "Failed to decode data from base64-encoded string: \(bytesBase64Encoded)"
+      )
+    }
+    self.init(mimeType: mimeType, data: data)
   }
 }

FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenGenerationResponseTests.swift

@@ -23,7 +23,9 @@ final class ImagenGenerationResponseTests: XCTestCase {
   func testDecodeResponse_oneBase64Image_noneFiltered() throws {
     let mimeType = "image/png"
     let bytesBase64Encoded = "dGVzdC1iYXNlNjQtZGF0YQ=="
-    let image = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded)
+    let image = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded))
+    )
     let json = """
     {
       "predictions": [
@@ -50,9 +52,15 @@ final class ImagenGenerationResponseTests: XCTestCase {
     let bytesBase64Encoded1 = "dGVzdC1iYXNlNjQtYnl0ZXMtMQ=="
     let bytesBase64Encoded2 = "dGVzdC1iYXNlNjQtYnl0ZXMtMg=="
     let bytesBase64Encoded3 = "dGVzdC1iYXNlNjQtYnl0ZXMtMw=="
-    let image1 = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded1)
-    let image2 = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded2)
-    let image3 = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded3)
+    let image1 = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded1))
+    )
+    let image2 = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded2))
+    )
+    let image3 = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded3))
+    )
     let json = """
     {
       "predictions": [
@@ -86,8 +94,12 @@ final class ImagenGenerationResponseTests: XCTestCase {
     let mimeType = "image/png"
     let bytesBase64Encoded1 = "dGVzdC1iYXNlNjQtYnl0ZXMtMQ=="
     let bytesBase64Encoded2 = "dGVzdC1iYXNlNjQtYnl0ZXMtMg=="
-    let image1 = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded1)
-    let image2 = ImagenInlineImage(mimeType: mimeType, bytesBase64Encoded: bytesBase64Encoded2)
+    let image1 = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded1))
+    )
+    let image2 = try ImagenInlineImage(
+      mimeType: mimeType, data: XCTUnwrap(Data(base64Encoded: bytesBase64Encoded2))
+    )
     let raiFilteredReason = """
     Your current safety filter threshold filtered out 2 generated images. You will not be charged \
     for blocked images. Try rephrasing the prompt. If you think this was an error, send feedback.

FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineImageTests.swift

@@ -77,4 +77,27 @@ final class ImagenInlineImageTests: XCTestCase {
       XCTFail("Expected a DecodingError.keyNotFound error; got \(error).")
     }
   }
+
+  func testDecodeImage_invalidBase64Data_throws() throws {
+    let bytesBase64Encoded = "not-a-base64-string"
+    let json = """
+    {
+      "bytesBase64Encoded": "\(bytesBase64Encoded)",
+      "mimeType": "image/png"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    do {
+      _ = try decoder.decode(ImagenInlineImage.self, from: jsonData)
+      XCTFail("Expected an error; none thrown.")
+    } catch let DecodingError.dataCorrupted(context) {
+      XCTAssertEqual(
+        context.debugDescription,
+        "Failed to decode data from base64-encoded string: \(bytesBase64Encoded)"
+      )
+    } catch {
+      XCTFail("Expected a DecodingError.dataCorrupted error; got \(error).")
+    }
+  }
 }