Skip to content

Commit 4362efa

Browse files
liustveliustve
authored
Add missing request/response attributes assertions in Bedrock AgentCore Contract Tests (#535)
*Description of changes:* Fixes a gap in Bedrock AgentCore contract tests to also validate the attributes in the request/response in BedrockAgentCore spans By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
1 parent e9842e6 commit 4362efa

File tree

1 file changed

+113
-4
lines changed

1 file changed

+113
-4
lines changed

‎contract-tests/tests/test/amazon/botocore/botocore_test.py‎

Lines changed: 113 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,11 @@
3939
_AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
4040
_AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
4141
_AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
42+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: str = "aws.bedrock.agentcore.runtime.arn"
43+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "aws.bedrock.agentcore.runtime_endpoint.arn"
44+
_AWS_BEDROCK_AGENTCORE_MEMORY_ARN: str = "aws.bedrock.agentcore.memory.arn"
45+
_AWS_BEDROCK_AGENTCORE_GATEWAY_ARN: str = "aws.bedrock.agentcore.gateway.arn"
46+
_AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "aws.bedrock.agentcore.identity.workload_identity.arn"
4247

4348
_GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"
4449
_GEN_AI_REQUEST_TEMPERATURE: str = "gen_ai.request.temperature"
@@ -48,7 +53,14 @@
4853
_GEN_AI_USAGE_INPUT_TOKENS: str = "gen_ai.usage.input_tokens"
4954
_GEN_AI_USAGE_OUTPUT_TOKENS: str = "gen_ai.usage.output_tokens"
5055
_GEN_AI_SYSTEM: str = "gen_ai.system"
51-
56+
_GEN_AI_RUNTIME_ID: str = "gen_ai.runtime.id"
57+
_GEN_AI_BROWSER_ID: str = "gen_ai.browser.id"
58+
_GEN_AI_CODE_INTERPRETER_ID: str = "gen_ai.code_interpreter.id"
59+
_GEN_AI_MEMORY_ID: str = "gen_ai.memory.id"
60+
_GEN_AI_GATEWAY_ID: str = "gen_ai.gateway.id"
61+
62+
_AWS_GATEWAY_TARGET_ID: str = "aws.gateway.target.id"
63+
_AWS_AUTH_CREDENTIAL_PROVIDER: str = "aws.auth.credential_provider"
5264
_AWS_SECRET_ARN: str = "aws.secretsmanager.secret.arn"
5365
_AWS_STATE_MACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
5466
_AWS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"
@@ -764,10 +776,20 @@ def test_bedrock_agentcore_create_agent_runtime(self):
764776
remote_resource_type="AWS::BedrockAgentCore::Runtime",
765777
remote_resource_identifier=expected_identifier,
766778
cloudformation_primary_identifier=expected_identifier,
779+
response_specific_attributes={
780+
_GEN_AI_RUNTIME_ID: expected_identifier,
781+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: (
782+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/{expected_identifier}"
783+
),
784+
_AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN: (
785+
"arn:aws:iam::123456789012:role/service-role/" "AmazonBedrockAgentCoreRuntimeDefaultServiceRole"
786+
),
787+
},
767788
span_name="Bedrock AgentCore Control.CreateAgentRuntime",
768789
)
769790

770791
def test_bedrock_agentcore_create_agent_runtime_endpoint(self):
792+
agent_id = "myAgent-w8slyU6q5M"
771793
self.do_test_requests(
772794
"bedrock-agentcore/runtime/createendpoint/myAgent-w8slyU6q5M",
773795
"GET",
@@ -782,6 +804,18 @@ def test_bedrock_agentcore_create_agent_runtime_endpoint(self):
782804
cloudformation_primary_identifier=(
783805
"arn:aws:bedrock-agentcore:us-west-2:123456789012:endpoint/invokeEndpoint"
784806
),
807+
request_specific_attributes={
808+
_GEN_AI_RUNTIME_ID: agent_id,
809+
},
810+
response_specific_attributes={
811+
_GEN_AI_RUNTIME_ID: agent_id,
812+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: (
813+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/{agent_id}"
814+
),
815+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN: (
816+
"arn:aws:bedrock-agentcore:us-west-2:123456789012:endpoint/invokeEndpoint"
817+
),
818+
},
785819
span_name="Bedrock AgentCore Control.CreateAgentRuntimeEndpoint",
786820
)
787821

@@ -799,6 +833,11 @@ def test_bedrock_agentcore_invoke_agent_runtime(self):
799833
remote_resource_type="AWS::BedrockAgentCore::Runtime",
800834
remote_resource_identifier=expected_identifier,
801835
cloudformation_primary_identifier=expected_identifier,
836+
request_specific_attributes={
837+
_AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: (
838+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/{expected_identifier}"
839+
),
840+
},
802841
span_name="Bedrock AgentCore.InvokeAgentRuntime",
803842
)
804843

@@ -816,6 +855,12 @@ def test_bedrock_agentcore_start_browser_session(self):
816855
remote_resource_type="AWS::BedrockAgentCore::BrowserCustom",
817856
remote_resource_identifier=expected_identifier,
818857
cloudformation_primary_identifier=expected_identifier,
858+
request_specific_attributes={
859+
_GEN_AI_BROWSER_ID: expected_identifier,
860+
},
861+
response_specific_attributes={
862+
_GEN_AI_BROWSER_ID: expected_identifier,
863+
},
819864
span_name="Bedrock AgentCore.StartBrowserSession",
820865
)
821866

@@ -833,6 +878,12 @@ def test_bedrock_agentcore_start_browser_session_v1(self):
833878
remote_resource_type="AWS::BedrockAgentCore::Browser",
834879
remote_resource_identifier=expected_identifier,
835880
cloudformation_primary_identifier=expected_identifier,
881+
request_specific_attributes={
882+
_GEN_AI_BROWSER_ID: expected_identifier,
883+
},
884+
response_specific_attributes={
885+
_GEN_AI_BROWSER_ID: expected_identifier,
886+
},
836887
span_name="Bedrock AgentCore.StartBrowserSession",
837888
)
838889

@@ -850,6 +901,12 @@ def test_bedrock_agentcore_start_code_interpreter_session(self):
850901
remote_resource_type="AWS::BedrockAgentCore::CodeInterpreterCustom",
851902
remote_resource_identifier=expected_identifier,
852903
cloudformation_primary_identifier=expected_identifier,
904+
request_specific_attributes={
905+
_GEN_AI_CODE_INTERPRETER_ID: expected_identifier,
906+
},
907+
response_specific_attributes={
908+
_GEN_AI_CODE_INTERPRETER_ID: expected_identifier,
909+
},
853910
span_name="Bedrock AgentCore.StartCodeInterpreterSession",
854911
)
855912

@@ -867,6 +924,12 @@ def test_bedrock_agentcore_start_code_interpreter_session_v1(self):
867924
remote_resource_type="AWS::BedrockAgentCore::CodeInterpreter",
868925
remote_resource_identifier=expected_identifier,
869926
cloudformation_primary_identifier=expected_identifier,
927+
request_specific_attributes={
928+
_GEN_AI_CODE_INTERPRETER_ID: expected_identifier,
929+
},
930+
response_specific_attributes={
931+
_GEN_AI_CODE_INTERPRETER_ID: expected_identifier,
932+
},
870933
span_name="Bedrock AgentCore.StartCodeInterpreterSession",
871934
)
872935

@@ -884,6 +947,12 @@ def test_bedrock_agentcore_create_event(self):
884947
remote_resource_type="AWS::BedrockAgentCore::Memory",
885948
remote_resource_identifier=expected_identifier,
886949
cloudformation_primary_identifier=expected_identifier,
950+
request_specific_attributes={
951+
_GEN_AI_MEMORY_ID: expected_identifier,
952+
},
953+
response_specific_attributes={
954+
_GEN_AI_MEMORY_ID: expected_identifier,
955+
},
887956
span_name="Bedrock AgentCore.CreateEvent",
888957
)
889958

@@ -903,6 +972,12 @@ def test_bedrock_agentcore_create_memory(self):
903972
cloudformation_primary_identifier=(
904973
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:memory/{expected_identifier}"
905974
),
975+
response_specific_attributes={
976+
_GEN_AI_MEMORY_ID: expected_identifier,
977+
_AWS_BEDROCK_AGENTCORE_MEMORY_ARN: (
978+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:memory/{expected_identifier}"
979+
),
980+
},
906981
span_name="Bedrock AgentCore Control.CreateMemory",
907982
)
908983

@@ -920,11 +995,18 @@ def test_bedrock_agentcore_create_gateway(self):
920995
remote_resource_type="AWS::BedrockAgentCore::Gateway",
921996
remote_resource_identifier=expected_identifier,
922997
cloudformation_primary_identifier=expected_identifier,
998+
response_specific_attributes={
999+
_GEN_AI_GATEWAY_ID: expected_identifier,
1000+
_AWS_BEDROCK_AGENTCORE_GATEWAY_ARN: (
1001+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:gateway/{expected_identifier}"
1002+
),
1003+
},
9231004
span_name="Bedrock AgentCore Control.CreateGateway",
9241005
)
9251006

9261007
def test_bedrock_agentcore_create_gateway_target(self):
927-
expected_identifier = "workinggateway-oersefsjga"
1008+
gateway_id = "workinggateway-oersefsjga"
1009+
target_id = "testTarget-123"
9281010
self.do_test_requests(
9291011
"bedrock-agentcore/gateway/creategatewaytarget/workinggateway-oersefsjga",
9301012
"GET",
@@ -935,8 +1017,17 @@ def test_bedrock_agentcore_create_gateway_target(self):
9351017
remote_service="AWS::BedrockAgentCore",
9361018
remote_operation="CreateGatewayTarget",
9371019
remote_resource_type="AWS::BedrockAgentCore::GatewayTarget",
938-
remote_resource_identifier=expected_identifier,
939-
cloudformation_primary_identifier=expected_identifier,
1020+
remote_resource_identifier=gateway_id,
1021+
cloudformation_primary_identifier=gateway_id,
1022+
request_specific_attributes={
1023+
_GEN_AI_GATEWAY_ID: gateway_id,
1024+
},
1025+
response_specific_attributes={
1026+
_AWS_GATEWAY_TARGET_ID: target_id,
1027+
_AWS_BEDROCK_AGENTCORE_GATEWAY_ARN: (
1028+
f"arn:aws:bedrock-agentcore:us-west-2:123456789012:gateway/{gateway_id}"
1029+
),
1030+
},
9401031
span_name="Bedrock AgentCore Control.CreateGatewayTarget",
9411032
)
9421033

@@ -954,6 +1045,12 @@ def test_bedrock_agentcore_create_oauth2_credential_provider(self):
9541045
remote_resource_type="AWS::BedrockAgentCore::OAuth2CredentialProvider",
9551046
remote_resource_identifier=expected_identifier,
9561047
cloudformation_primary_identifier=expected_identifier,
1048+
response_specific_attributes={
1049+
_AWS_AUTH_CREDENTIAL_PROVIDER: (
1050+
f"arn:aws:acps:us-west-2:123456789012:"
1051+
f"token-vault/default/oauth2credentialprovider/{expected_identifier}"
1052+
),
1053+
},
9571054
span_name="Bedrock AgentCore Control.CreateOauth2CredentialProvider",
9581055
)
9591056

@@ -971,6 +1068,12 @@ def test_bedrock_agentcore_create_api_key_credential_provider(self):
9711068
remote_resource_type="AWS::BedrockAgentCore::APIKeyCredentialProvider",
9721069
remote_resource_identifier=expected_identifier,
9731070
cloudformation_primary_identifier=expected_identifier,
1071+
response_specific_attributes={
1072+
_AWS_AUTH_CREDENTIAL_PROVIDER: (
1073+
f"arn:aws:acps:us-west-2:123456789012:"
1074+
f"token-vault/default/apikeycredentialprovider/{expected_identifier}"
1075+
),
1076+
},
9741077
span_name="Bedrock AgentCore Control.CreateApiKeyCredentialProvider",
9751078
)
9761079

@@ -988,6 +1091,9 @@ def test_bedrock_agentcore_get_resource_oauth2_token(self):
9881091
remote_resource_type="AWS::BedrockAgentCore::OAuth2CredentialProvider",
9891092
remote_resource_identifier=expected_identifier,
9901093
cloudformation_primary_identifier=expected_identifier,
1094+
request_specific_attributes={
1095+
_AWS_AUTH_CREDENTIAL_PROVIDER: expected_identifier,
1096+
},
9911097
span_name="Bedrock AgentCore.GetResourceOauth2Token",
9921098
)
9931099

@@ -1005,6 +1111,9 @@ def test_bedrock_agentcore_get_resource_api_key(self):
10051111
remote_resource_type="AWS::BedrockAgentCore::APIKeyCredentialProvider",
10061112
remote_resource_identifier=expected_identifier,
10071113
cloudformation_primary_identifier=expected_identifier,
1114+
request_specific_attributes={
1115+
_AWS_AUTH_CREDENTIAL_PROVIDER: expected_identifier,
1116+
},
10081117
span_name="Bedrock AgentCore.GetResourceApiKey",
10091118
)
10101119

0 commit comments

Comments
 (0)