test: improve

Author: leotac

tests_integ/test_bedrock_guardrails.py

@@ -205,12 +205,15 @@ def test_guardrail_output_intervention_redact_output(bedrock_guardrail, processi
 
 @pytest.mark.parametrize("processing_mode", ["sync", "async"])
 def test_guardrail_intervention_properly_redacts_tool_result(bedrock_guardrail, processing_mode):
-    REDACT_MESSAGE = "Input redacted."
+    INPUT_REDACT_MESSAGE = "Input redacted."
+    OUTPUT_REDACT_MESSAGE = "Output redacted."
     bedrock_model = BedrockModel(
         guardrail_id=bedrock_guardrail,
         guardrail_version="DRAFT",
         guardrail_stream_processing_mode=processing_mode,
-        guardrail_redact_input_message=REDACT_MESSAGE,
+        guardrail_redact_output=True,
+        guardrail_redact_input_message=INPUT_REDACT_MESSAGE,
+        guardrail_redact_output_message=OUTPUT_REDACT_MESSAGE,
         region_name="us-east-1",
     )
 
@@ -230,51 +233,44 @@ def list_users() -> str:
     response1 = agent("List my users.")
     response2 = agent("Thank you!")
 
-    assert response1.stop_reason == "guardrail_intervened"
-
-    """
-    In async streaming: The buffering is non-blocking. 
-    Tokens are streamed while Guardrails processes the buffered content in the background. 
-    This means the response may be returned before Guardrails has finished processing.
-    As a result, we cannot guarantee that the BLOCKED_OUTPUT is in the response
-    However, response2 should not be blocked anyway.
+    """ Message sequence:
+    0 (user): request1
+    1 (assistant): reasoning + tool call
+    2 (user): tool result
+    3 (assistant): response1 -> output guardrail intervenes
+    4 (user): request2
+    5 (assistant): response2
+    
+    Guardrail intervened on output in message 3 will cause
+    the redaction of the preceding input (message 2) and message 3.
     """
+
     if processing_mode == "sync":
-        assert response1.stop_reason == "guardrail_intervened"
-        assert BLOCKED_OUTPUT in str(response1)
+        """ In sync mode the guardrail processing is blocking.
+        The response is already blocked and redacted. """
 
-        assert response2.stop_reason != "guardrail_intervened"
-        assert BLOCKED_OUTPUT not in str(response2)
+        assert OUTPUT_REDACT_MESSAGE in str(response1)
+        assert OUTPUT_REDACT_MESSAGE not in str(response2)
+
+    """
+    In async streaming, the buffering is non-blocking,
+    so the response may be returned before Guardrails has finished processing.
+    
+    However, in both sync and async, with guardrail_redact_output=True:
+    1. the content should be properly redacted in memory allowing the 
+    conversation to continue;
+    """
+    assert response1.stop_reason == "guardrail_intervened"
+    assert response2.stop_reason != "guardrail_intervened"
 
-        """ Message sequence:
-        0 (user): request1
-        1 (assistant): reasoning + tool call
-        2 (user): tool result
-        3 (assistant): response1 (blocked)
-        4 (user): request2
-        5 (assistant): response2
-        
-        Guardrail intervened on output in message 3 will cause
-        the redaction of the preceding input (message 2).
-        We want the tool result block to be preserved.
-        """
-
-        tool_call = [b for b in agent.messages[1]["content"] if "toolUse" in b][0]
-        tool_result = [b for b in agent.messages[2]["content"] if "toolResult" in b][0]
-        assert tool_result["tool_id"] == tool_call["tool_id"]
-        assert tool_result["content"][0]["text"] == REDACT_MESSAGE
+    """
+    2. the tool result block should be redacted properly.
+    """
 
-    else:
-        # TODO
-        cactus_returned_in_response1_blocked_by_input_guardrail = BLOCKED_INPUT in str(response2)
-        cactus_blocked_in_response1_allows_next_response = (
-            REDACT_MESSAGE not in str(response2) and response2.stop_reason != "guardrail_intervened"
-        )
-        assert (
-            cactus_returned_in_response1_blocked_by_input_guardrail or cactus_blocked_in_response1_allows_next_response
-        )
-        # Output correctly redacted
-        assert agent.messages[3]["content"][0]["text"] == REDACT_MESSAGE
+    tool_call = [b for b in agent.messages[1]["content"] if "toolUse" in b][0]["toolUse"]
+    tool_result = [b for b in agent.messages[2]["content"] if "toolResult" in b][0]["toolResult"]
+    assert tool_result["toolUseId"] == tool_call["toolUseId"]
+    assert tool_result["content"][0]["text"] == INPUT_REDACT_MESSAGE
 
 
 def test_guardrail_input_intervention_properly_redacts_in_session(boto_session, bedrock_guardrail, temp_dir):