Add contract tests

Author: blairhyy-amazon

contract-tests/images/applications/botocore/botocore_server.py

@@ -54,6 +54,8 @@ def do_GET(self):
             self._handle_stepfunctions_request()
         if self.in_path("sns"):
             self._handle_sns_request()
+        if self.in_path("cross-account"):
+            self._handle_cross_account_request()
 
         self._end_request(self.main_status)
 
@@ -85,6 +87,23 @@ def do_PUT(self):
     def in_path(self, sub_path: str) -> bool:
         return sub_path in self.path
 
+    def _handle_cross_account_request(self) -> None:
+        s3_client = boto3.client(
+            "s3",
+            endpoint_url=_AWS_SDK_S3_ENDPOINT,
+            region_name="eu-central-1",
+            aws_access_key_id="account_b_access_key_id",
+            aws_secret_access_key="account_b_secret_access_key",
+            aws_session_token="account_b_token",
+        )
+        if self.in_path("createbucket/account_b"):
+            set_main_status(200)
+            s3_client.create_bucket(
+                Bucket="cross-account-bucket", CreateBucketConfiguration={"LocationConstraint": "eu-central-1"}
+            )
+        else:
+            set_main_status(404)
+
     def _handle_s3_request(self) -> None:
         s3_client: BaseClient = boto3.client("s3", endpoint_url=_AWS_SDK_S3_ENDPOINT, region_name=_AWS_REGION)
         if self.in_path(_ERROR):
@@ -151,6 +170,11 @@ def _handle_ddb_request(self) -> None:
                 ],
                 BillingMode="PAY_PER_REQUEST",
             )
+        elif self.in_path("describetable/some-table"):
+            set_main_status(200)
+            ddb_client.describe_table(
+                TableName="put_test_table",
+            )
         elif self.in_path("putitem/putitem-table/key"):
             set_main_status(200)
             item: dict = {"id": {"S": "1"}}
@@ -213,6 +237,11 @@ def _handle_kinesis_request(self) -> None:
         elif self.in_path("putrecord/my-stream"):
             set_main_status(200)
             kinesis_client.put_record(StreamName="test_stream", Data=b"test", PartitionKey="partition_key")
+        elif self.in_path("describestream/my-stream"):
+            set_main_status(200)
+            kinesis_client.describe_stream(
+                StreamName="test_stream", StreamARN="arn:aws:kinesis:us-west-2:000000000000:stream/test_stream"
+            )
         else:
             set_main_status(404)
 

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

@@ -16,7 +16,10 @@
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_OPERATION,
+    AWS_REMOTE_RESOURCE_ACCESS_KEY,
+    AWS_REMOTE_RESOURCE_ACCOUNT_ID,
     AWS_REMOTE_RESOURCE_IDENTIFIER,
+    AWS_REMOTE_RESOURCE_REGION,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
@@ -50,6 +53,8 @@
 _AWS_STATE_MACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
 _AWS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"
 _AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
+_AWS_DYNAMODB_TABLE_ARN: str = "aws.dynamodb.table.arn"
+_AWS_KINESIS_STREAM_ARN: str = "aws.kinesis.stream.arn"
 
 
 # pylint: disable=too-many-public-methods,too-many-lines
@@ -212,6 +217,29 @@ def test_dynamodb_create_table(self):
             span_name="DynamoDB.CreateTable",
         )
 
+    def test_dynamodb_describe_table(self):
+        self.do_test_requests(
+            "ddb/describetable/some-table",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::DynamoDB",
+            remote_operation="DescribeTable",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="put_test_table",
+            remote_resource_account_id="000000000000",
+            remote_resource_region="us-west-2",
+            cloudformation_primary_identifier="put_test_table",
+            request_specific_attributes={
+                SpanAttributes.AWS_DYNAMODB_TABLE_NAMES: ["put_test_table"],
+            },
+            response_specific_attributes={
+                _AWS_DYNAMODB_TABLE_ARN: r"arn:aws:dynamodb:us-west-2:000000000000:table/put_test_table",
+            },
+            span_name="DynamoDB.DescribeTable",
+        )
+
     def test_dynamodb_put_item(self):
         self.do_test_requests(
             "ddb/putitem/putitem-table/key",
@@ -379,6 +407,26 @@ def test_kinesis_put_record(self):
             span_name="Kinesis.PutRecord",
         )
 
+    def test_kinesis_describe_stream(self):
+        self.do_test_requests(
+            "kinesis/describestream/my-stream",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::Kinesis",
+            remote_operation="DescribeStream",
+            remote_resource_type="AWS::Kinesis::Stream",
+            remote_resource_identifier="test_stream",
+            cloudformation_primary_identifier="test_stream",
+            remote_resource_account_id="000000000000",
+            remote_resource_region="us-west-2",
+            request_specific_attributes={
+                _AWS_KINESIS_STREAM_NAME: "test_stream",
+            },
+            span_name="Kinesis.DescribeStream",
+        )
+
     def test_kinesis_error(self):
         self.do_test_requests(
             "kinesis/error",
@@ -905,6 +953,26 @@ def test_stepfunctions_fault(self):
             span_name="SFN.ListStateMachineVersions",
         )
 
+    def test_cross_account(self):
+        self.do_test_requests(
+            "cross-account/createbucket/account_b",
+            "GET",
+            200,
+            0,
+            0,
+            remote_service="AWS::S3",
+            remote_operation="CreateBucket",
+            remote_resource_type="AWS::S3::Bucket",
+            remote_resource_identifier="cross-account-bucket",
+            cloudformation_primary_identifier="cross-account-bucket",
+            request_specific_attributes={
+                SpanAttributes.AWS_S3_BUCKET: "cross-account-bucket",
+            },
+            remote_resource_access_key="account_b_access_key_id",
+            remote_resource_region="eu-central-1",
+            span_name="S3.CreateBucket",
+        )
+
     # TODO: Add contract test for lambda event source mapping resource
 
     @override
@@ -924,6 +992,9 @@ def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSp
             kwargs.get("remote_resource_type", "None"),
             kwargs.get("remote_resource_identifier", "None"),
             kwargs.get("cloudformation_primary_identifier", "None"),
+            kwargs.get("remote_resource_account_id", "None"),
+            kwargs.get("remote_resource_access_key", "None"),
+            kwargs.get("remote_resource_region", "None"),
         )
 
     def _assert_aws_attributes(
@@ -935,6 +1006,9 @@ def _assert_aws_attributes(
         remote_resource_type: str,
         remote_resource_identifier: str,
         cloudformation_primary_identifier: str,
+        remote_resource_account_id: str,
+        remote_resource_access_key: str,
+        remote_resource_region: str,
     ) -> None:
         attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
         self._assert_str_attribute(attributes_dict, AWS_LOCAL_SERVICE, self.get_application_otel_service_name())
@@ -961,8 +1035,16 @@ def _assert_aws_attributes(
                 self._assert_str_attribute(
                     attributes_dict, AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, cloudformation_primary_identifier
                 )
-        # See comment above AWS_LOCAL_OPERATION
-        self._assert_str_attribute(attributes_dict, AWS_SPAN_KIND, span_kind)
+        if remote_resource_account_id != "None":
+            assert remote_resource_identifier != "None"
+            self._assert_str_attribute(attributes_dict, AWS_REMOTE_RESOURCE_ACCOUNT_ID, remote_resource_account_id)
+            self.assertIsNone(attributes_dict.get(AWS_REMOTE_RESOURCE_ACCESS_KEY))
+        if remote_resource_access_key != "None":
+            assert remote_resource_identifier != "None"
+            self._assert_str_attribute(attributes_dict, AWS_REMOTE_RESOURCE_ACCESS_KEY, remote_resource_access_key)
+            self.assertIsNone(attributes_dict.get(AWS_REMOTE_RESOURCE_ACCOUNT_ID))
+        if remote_resource_region != "None":
+            self._assert_str_attribute(attributes_dict, AWS_REMOTE_RESOURCE_REGION, remote_resource_region)
 
     @override
     def _assert_semantic_conventions_span_attributes(
@@ -1050,13 +1132,26 @@ def _assert_metric_attributes(
         self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, "CLIENT")
         remote_resource_type = kwargs.get("remote_resource_type", "None")
         remote_resource_identifier = kwargs.get("remote_resource_identifier", "None")
+        remote_resource_account_id = kwargs.get("remote_resource_account_id", "None")
+        remote_resource_access_key = kwargs.get("remote_resource_access_key", "None")
+        remote_resource_region = kwargs.get("remote_resource_region", "None")
         if remote_resource_type != "None":
             self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_TYPE, remote_resource_type)
         if remote_resource_identifier != "None":
             if self._is_valid_regex(remote_resource_identifier):
                 self._assert_match_attribute(attribute_dict, AWS_REMOTE_RESOURCE_IDENTIFIER, remote_resource_identifier)
             else:
                 self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_IDENTIFIER, remote_resource_identifier)
+        if remote_resource_account_id != "None":
+            assert remote_resource_identifier != "None"
+            self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_ACCOUNT_ID, remote_resource_account_id)
+            self.assertIsNone(attribute_dict.get(AWS_REMOTE_RESOURCE_ACCESS_KEY))
+        if remote_resource_access_key != "None":
+            assert remote_resource_identifier != "None"
+            self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_ACCESS_KEY, remote_resource_access_key)
+            self.assertIsNone(attribute_dict.get(AWS_REMOTE_RESOURCE_ACCOUNT_ID))
+        if remote_resource_region != "None":
+            self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_REGION, remote_resource_region)
         self.check_sum(metric_name, dependency_dp.sum, expected_sum)
 
         attribute_dict: Dict[str, AnyValue] = self._get_attributes_dict(service_dp.attributes)

contract-tests/tests/test/amazon/utils/application_signals_constants.py

@@ -22,6 +22,9 @@
 AWS_REMOTE_DB_USER: str = "aws.remote.db.user"
 AWS_REMOTE_SERVICE: str = "aws.remote.service"
 AWS_REMOTE_OPERATION: str = "aws.remote.operation"
+AWS_REMOTE_RESOURCE_ACCESS_KEY: str = "aws.remote.resource.account.access_key"
+AWS_REMOTE_RESOURCE_ACCOUNT_ID: str = "aws.remote.resource.account.id"
+AWS_REMOTE_RESOURCE_REGION: str = "aws.remote.resource.region"
 AWS_REMOTE_RESOURCE_TYPE: str = "aws.remote.resource.type"
 AWS_REMOTE_RESOURCE_IDENTIFIER: str = "aws.remote.resource.identifier"
 AWS_SPAN_KIND: str = "aws.span.kind"