Skip to content

Commit f807fe7

Browse files
committed
Fix filteredReasons concatenation and update tests
1 parent e18f8bc commit f807fe7

File tree

2 files changed

+19
-11
lines changed

2 files changed

+19
-11
lines changed

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

+6-1
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,12 @@ extension ImagenGenerationResponse: Decodable where T: Decodable {
7171
}
7272

7373
self.images = images
74-
filteredReason = filteredReasons.joined(separator: "\n")
74+
let filteredReason = filteredReasons.joined(separator: "\n")
75+
if filteredReason.isEmpty {
76+
self.filteredReason = nil
77+
} else {
78+
self.filteredReason = filteredReason
79+
}
7580
// TODO(#14221): Throw `ImagenImagesBlockedError` with `filteredReason` if `images` is empty.
7681
}
7782
}

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

+13-10
Original file line numberDiff line numberDiff line change
@@ -181,16 +181,20 @@ final class ImagenGenerationResponseTests: XCTestCase {
181181
let json = "{}"
182182
let jsonData = try XCTUnwrap(json.data(using: .utf8))
183183

184-
let response = try decoder.decode(
185-
ImagenGenerationResponse<ImagenInlineImage>.self,
186-
from: jsonData
187-
)
188-
189-
XCTAssertEqual(response.images, [])
190-
XCTAssertNil(response.filteredReason)
184+
do {
185+
let response = try decoder.decode(
186+
ImagenGenerationResponse<ImagenInlineImage>.self,
187+
from: jsonData
188+
)
189+
XCTFail("Expected a DecodingError, got response: \(response)")
190+
} catch let DecodingError.keyNotFound(codingKey, _) {
191+
XCTAssertEqual(codingKey.stringValue, "predictions")
192+
} catch {
193+
XCTFail("Expected a DecodingError.keyNotFound, got error: \(error)")
194+
}
191195
}
192196

193-
func testDecodeResponse_multipleFilterReasons_returnsFirst() throws {
197+
func testDecodeResponse_multipleFilterReasons_concatenatesReasons() throws {
194198
let raiFilteredReason1 = "filtered-reason-1"
195199
let raiFilteredReason2 = "filtered-reason-2"
196200
let json = """
@@ -213,8 +217,7 @@ final class ImagenGenerationResponseTests: XCTestCase {
213217
)
214218

215219
XCTAssertEqual(response.images, [])
216-
XCTAssertEqual(response.filteredReason, raiFilteredReason1)
217-
XCTAssertNotEqual(response.filteredReason, raiFilteredReason2)
220+
XCTAssertEqual(response.filteredReason, "\(raiFilteredReason1)\n\(raiFilteredReason2)")
218221
}
219222

220223
func testDecodeResponse_unknownPrediction() throws {

0 commit comments

Comments
 (0)