From f4d81ac48b3d03066535ce8fc0380f42cbe310ab Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Fri, 14 Feb 2025 15:49:32 -0500 Subject: [PATCH] [Vertex AI] Add error handling for decoding `ImagenInlineImage` from b64 --- .../Public/Imagen/ImagenInlineImage.swift | 15 +++++++----- .../ImagenGenerationResponseTests.swift | 24 ++++++++++++++----- .../Types/Imagen/ImagenInlineImageTests.swift | 23 ++++++++++++++++++ 3 files changed, 50 insertions(+), 12 deletions(-) diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineImage.swift index 0f7f030a218..a526dca3e56 100644 --- a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineImage.swift +++ b/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) } } diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenGenerationResponseTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenGenerationResponseTests.swift index e52db094028..4ea2b8a150e 100644 --- a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenGenerationResponseTests.swift +++ b/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. diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineImageTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineImageTests.swift index fa8072ee497..a28bf763b1b 100644 --- a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineImageTests.swift +++ b/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).") + } + } }