@@ -133,6 +133,20 @@ def mock_generate_content(
133
133
)
134
134
return response
135
135
136
+ should_block = (
137
+ last_message_part .text
138
+ and "Please block with block_reason=OTHER" in last_message_part .text
139
+ )
140
+ if should_block :
141
+ response = gapic_prediction_service_types .GenerateContentResponse (
142
+ candidates = [],
143
+ prompt_feedback = gapic_prediction_service_types .GenerateContentResponse .PromptFeedback (
144
+ block_reason = gapic_prediction_service_types .GenerateContentResponse .PromptFeedback .BlockedReason .OTHER ,
145
+ block_reason_message = "Blocked for testing" ,
146
+ ),
147
+ )
148
+ return response
149
+
136
150
is_continued_chat = len (request .contents ) > 1
137
151
has_retrieval = any (
138
152
tool .retrieval or tool .google_search_retrieval for tool in request .tools
@@ -349,6 +363,31 @@ def test_chat_send_message_response_validation_errors(
349
363
# Checking that history did not get updated
350
364
assert len (chat .history ) == 2
351
365
366
+ @mock .patch .object (
367
+ target = prediction_service .PredictionServiceClient ,
368
+ attribute = "generate_content" ,
369
+ new = mock_generate_content ,
370
+ )
371
+ @pytest .mark .parametrize (
372
+ "generative_models" ,
373
+ [generative_models , preview_generative_models ],
374
+ )
375
+ def test_chat_send_message_response_blocked_errors (
376
+ self , generative_models : generative_models
377
+ ):
378
+ model = generative_models .GenerativeModel ("gemini-pro" )
379
+ chat = model .start_chat ()
380
+ response1 = chat .send_message ("Why is sky blue?" )
381
+ assert response1 .text
382
+ assert len (chat .history ) == 2
383
+
384
+ with pytest .raises (generative_models .ResponseValidationError ) as e :
385
+ chat .send_message ("Please block with block_reason=OTHER." )
386
+
387
+ assert e .match ("Blocked for testing" )
388
+ # Checking that history did not get updated
389
+ assert len (chat .history ) == 2
390
+
352
391
@mock .patch .object (
353
392
target = prediction_service .PredictionServiceClient ,
354
393
attribute = "generate_content" ,
0 commit comments