@@ -619,6 +619,29 @@ final class GenerativeModelTests: XCTestCase {
619
619
XCTFail ( " Should throw " )
620
620
} catch let GenerateContentError . promptBlocked( response) {
621
621
XCTAssertNil ( response. text)
622
+ let promptFeedback = try XCTUnwrap ( response. promptFeedback)
623
+ XCTAssertEqual ( promptFeedback. blockReason, PromptFeedback . BlockReason. safety)
624
+ XCTAssertNil ( promptFeedback. blockReasonMessage)
625
+ } catch {
626
+ XCTFail ( " Should throw a promptBlocked " )
627
+ }
628
+ }
629
+
630
+ func testGenerateContent_failure_promptBlockedSafetyWithMessage( ) async throws {
631
+ MockURLProtocol
632
+ . requestHandler = try httpRequestHandler (
633
+ forResource: " unary-failure-prompt-blocked-safety-with-message " ,
634
+ withExtension: " json "
635
+ )
636
+
637
+ do {
638
+ _ = try await model. generateContent ( testPrompt)
639
+ XCTFail ( " Should throw " )
640
+ } catch let GenerateContentError . promptBlocked( response) {
641
+ XCTAssertNil ( response. text)
642
+ let promptFeedback = try XCTUnwrap ( response. promptFeedback)
643
+ XCTAssertEqual ( promptFeedback. blockReason, PromptFeedback . BlockReason. safety)
644
+ XCTAssertEqual ( promptFeedback. blockReasonMessage, " Reasons " )
622
645
} catch {
623
646
XCTFail ( " Should throw a promptBlocked " )
624
647
}
@@ -909,7 +932,31 @@ final class GenerativeModelTests: XCTestCase {
909
932
XCTFail ( " Content shouldn't be shown, this shouldn't happen. " )
910
933
}
911
934
} catch let GenerateContentError . promptBlocked( response) {
912
- XCTAssertEqual ( response. promptFeedback? . blockReason, . safety)
935
+ let promptFeedback = try XCTUnwrap ( response. promptFeedback)
936
+ XCTAssertEqual ( promptFeedback. blockReason, . safety)
937
+ XCTAssertNil ( promptFeedback. blockReasonMessage)
938
+ return
939
+ }
940
+
941
+ XCTFail ( " Should have caught an error. " )
942
+ }
943
+
944
+ func testGenerateContentStream_failurePromptBlockedSafetyWithMessage( ) async throws {
945
+ MockURLProtocol
946
+ . requestHandler = try httpRequestHandler (
947
+ forResource: " streaming-failure-prompt-blocked-safety-with-message " ,
948
+ withExtension: " txt "
949
+ )
950
+
951
+ do {
952
+ let stream = try model. generateContentStream ( " Hi " )
953
+ for try await _ in stream {
954
+ XCTFail ( " Content shouldn't be shown, this shouldn't happen. " )
955
+ }
956
+ } catch let GenerateContentError . promptBlocked( response) {
957
+ let promptFeedback = try XCTUnwrap ( response. promptFeedback)
958
+ XCTAssertEqual ( promptFeedback. blockReason, . safety)
959
+ XCTAssertEqual ( promptFeedback. blockReasonMessage, " Reasons " )
913
960
return
914
961
}
915
962
0 commit comments