Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
Expand Down Expand Up @@ -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)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand All @@ -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": [
Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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).")
}
}
}
Loading