8000 chore(tests): Add unit tests for user agent changes (#125) · strands-agents/sdk-python@e907ac6 · GitHub
[go: up one dir, main page]
Skip to content

Commit e907ac6

Browse files
clareliguoriclareliguori
authored
chore(tests): Add unit tests for user agent changes (#125)
1 parent b04c4be commit e907ac6

File tree

1 file changed

+97
-8
lines changed

1 file changed

+97
-8
lines changed

‎tests/strands/models/test_bedrock.py

Lines changed: 97 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import boto3
55
import pytest
6+
from botocore.config import Config as BotocoreConfig
67
from botocore.exceptions import ClientError, EventStreamError
78

89
import strands
@@ -116,6 +117,54 @@ def test__init__with_region_and_session_raises_value_error():
116117
_ = BedrockModel(region_name="us-east-1", boto_session=boto3.Session(region_name="us-east-1"))
117118

118119

120+
def test__init__default_user_agent(bedrock_client):
121+
"""Set user agent when no boto_client_config is provided."""
122+
with unittest.mock.patch("strands.models.bedrock.boto3.Session") as mock_session_cls:
123+
mock_session = mock_session_cls.return_value
124+
_ = BedrockModel()
125+
126+
# Verify the client was created with the correct config
127+
mock_session.client.assert_called_once()
128+
args, kwargs = mock_session.client.call_args
129+
assert kwargs["service_name"] == "bedrock-runtime"
130+
assert isinstance(kwargs["config"], BotocoreConfig)
131+
assert kwargs["config"].user_agent_extra == "strands-agents"
132+
133+
134+
def test__init__with_custom_boto_client_config_no_user_agent(bedrock_client):
135+
"""Set user agent when boto_client_config is provided without user_agent_extra."""
136+
custom_config = BotocoreConfig(read_timeout=900)
137+
138+
with unittest.mock.patch("strands.models.bedrock.boto3.Session") as mock_session_cls:
139+
mock_session = mock_session_cls.return_value
140+
_ = BedrockModel(boto_client_config=custom_config)
141+
142+
# Verify the client was created with the correct config
143+
mock_session.client.assert_called_once()
144+
args, kwargs = mock_session.client.call_args
145+
assert kwargs["service_name"] == "bedrock-runtime"
146+
assert isinstance(kwargs["config"], BotocoreConfig)
147+
assert kwargs["config"].user_agent_extra == "strands-agents"
148+
assert kwargs["config"].read_timeout == 900
149+
150+
151+
def test__init__with_custom_boto_client_config_with_user_agent(bedrock_client):
152+
"""Append to existing user agent when boto_client_config is provided with user_agent_extra."""
153+
custom_config = BotocoreConfig(user_agent_extra="existing-agent", read_timeout=900)
154+
155+
with unittest.mock.patch("strands.models.bedrock.boto3.Session") as mock_session_cls:
156+
mock_session = mock_session_cls.return_value
157+
_ = BedrockModel(boto_client_config=custom_config)
158+
159+
# Verify the client was created with the correct config
160+
mock_session.client.assert_called_once()
161+
args, kwargs = mock_session.client.call_args
162+
assert kwargs["service_name"] == "bedrock-runtime"
163+
assert isinstance(kwargs["config"], BotocoreConfig)
164+
assert kwargs["config"].user_agent_extra == "existing-agent strands-agents"
165+
assert kwargs["config"].read_timeout == 900
166+
167+
119168
def test__init__model_config(bedrock_client):
120169
_ = bedrock_client
121170

@@ -381,7 +430,15 @@ def test_converse_input_guardrails(bedrock_client, model, messages, tool_spec, m
381430
"guardrail": {
382431
"inputAssessment": {
383432
"3e59qlue4hag": {
384-
"wordPolicy": {"customWords": [{"match": "CACTUS", "action": "BLOCKED", "detected": True}]}
433+
"wordPolicy": {
434+
"customWords": [
435+
{
436+
"match": "CACTUS",
437+
"action": "BLOCKED",
438+
"detected": True,
439+
}
440+
]
441+
}
385442
}
386443
}
387444
}
@@ -406,7 +463,10 @@ def test_converse_input_guardrails(bedrock_client, model, messages, tool_spec, m
406463
chunks = model.converse(messages, [tool_spec])
407464

408465
tru_chunks = list(chunks)
409-
exp_chunks = [{"redactContent": {"redactUserContentMessage": "[User input redacted.]"}}, metadata_event]
466+
exp_chunks = [
467+
{"redactContent": {"redactUserContentMessage": "[User input redacted.]"}},
468+
metadata_event,
469+
]
410470

411471
assert tru_chunks == exp_chunks
412472
bedrock_client.converse_stream.assert_called_once_with(**request)
@@ -424,7 +484,13 @@ def test_converse_output_guardrails(bedrock_client, model, messages, tool_spec,
424484
"3e59qlue4hag": [
425485
{
426486
"wordPolicy": {
427-
"customWords": [{"match": "CACTUS", "action": "BLOCKED", "detected": True}]
487+
"customWords": [
488+
{
489+
"match": "CACTUS",
490+
"action": "BLOCKED",
491+
"detected": True,
492+
}
493+
]
428494
},
429495
}
430496
]
@@ -451,7 +517,10 @@ def test_converse_output_guardrails(bedrock_client, model, messages, tool_spec,
451517
chunks = model.converse(messages, [tool_spec])
452518

453519
tru_chunks = list(chunks)
454-
exp_chunks = [{"redactContent": {"redactAssistantContentMessage": "[Assistant output redacted.]"}}, metadata_event]
520+
exp_chunks = [
521+
{"redactContent": {"redactAssistantContentMessage": "[Assistant output redacted.]"}},
522+
metadata_event,
523+
]
455524

456525
assert tru_chunks == exp_chunks
457526
bedrock_client.converse_stream.assert_called_once_with(**request)
@@ -471,7 +540,13 @@ def test_converse_output_guardrails_redacts_input_and_output(
471540
"3e59qlue4hag": [
472541
{
473542
"wordPolicy": {
474-
"customWords": [{"match": "CACTUS", "action": "BLOCKED", "detected": True}]
543+
"customWords": [
544+
{
545+
"match": "CACTUS",
546+
"action": "BLOCKED",
547+
"detected": True,
548+
}
549+
]
475550
},
476551
}
477552
]
@@ -521,7 +596,13 @@ def test_converse_output_no_blocked_guardrails_doesnt_redact(
521596
"3e59qlue4hag": [
522597
{
523598
"wordPolicy": {
524-
"customWords": [{"match": "CACTUS", "action": "NONE", "detected": True}]
599+
"customWords": [
600+
{
601+
"match": "CACTUS",
602+
"action": "NONE",
603+
"detected": True,
604+
}
605+
]
525606
},
526607
}
527608
]
@@ -567,7 +648,13 @@ def test_converse_output_no_guardrail_redact(
567648
"3e59qlue4hag": [
568649
{
569650
"wordPolicy": {
570-
"customWords": [{"match": "CACTUS", "action": "BLOCKED", "detected": True}]
651+
"customWords": [
652+
{
653+
"match": "CACTUS",
654+
"action": "BLOCKED",
655+
"detected": True,
656+
}
657+
]
571658
},
572659
}
573660
]
@@ -591,7 +678,9 @@ def test_converse_output_no_guardrail_redact(
591678
}
592679

593680
model.update_config(
594-
additional_request_fields=additional_request_fields, guardrail_redact_output=False, guardrail_redact_input=False
681+
additional_request_fields=additional_request_fields,
682+
guardrail_redact_output=False,
683+
guardrail_redact_input=False,
595684
)
596685
chunks = model.converse(messages, [tool_spec])
597686

0 commit comments

Comments
 (0)
0