Add contract tests

Author: blairhyy-amazon

contract-tests/images/applications/aws-sdk/server.js

@@ -9,9 +9,9 @@ const fetch = require('node-fetch');
 const JSZip = require('jszip');
 
 const { S3Client, CreateBucketCommand, PutObjectCommand, GetObjectCommand } = require('@aws-sdk/client-s3');
-const { DynamoDBClient, CreateTableCommand, PutItemCommand } = require('@aws-sdk/client-dynamodb');
+const { DynamoDBClient, CreateTableCommand, PutItemCommand, DescribeTableCommand } = require('@aws-sdk/client-dynamodb');
 const { SQSClient, CreateQueueCommand, SendMessageCommand, ReceiveMessageCommand } = require('@aws-sdk/client-sqs');
-const { KinesisClient, CreateStreamCommand, PutRecordCommand } = require('@aws-sdk/client-kinesis');
+const { KinesisClient, CreateStreamCommand, PutRecordCommand, DescribeStreamCommand } = require('@aws-sdk/client-kinesis');
 const { BedrockClient, GetGuardrailCommand } = require('@aws-sdk/client-bedrock');
 const { BedrockAgentClient, GetKnowledgeBaseCommand, GetDataSourceCommand, GetAgentCommand } = require('@aws-sdk/client-bedrock-agent');
 const { BedrockRuntimeClient, InvokeModelCommand } = require('@aws-sdk/client-bedrock-runtime');
@@ -194,6 +194,8 @@ async function handleGetRequest(req, res, path) {
     await handleSnsRequest(req, res, path);
   } else if (path.includes('lambda')) {
     await handleLambdaRequest(req, res, path);
+  } else if (path.includes('cross-account')) {
+    await handleCrossAccountRequest(req, res, path);
   } else {
     res.writeHead(404);
     res.end();
@@ -383,6 +385,19 @@ async function handleDdbRequest(req, res, path) {
       res.statusCode = 500;
     }
     res.end();
+  } else if (path.includes('describe/some-table')) {
+    try {
+      const response = await ddbClient.send(
+          new DescribeTableCommand({
+            TableName: 'put_test_table'
+          })
+        );
+      res.statusCode = 200;
+    } catch (err) {
+      console.log('Error describing the table', err);
+      res.statusCode = 500;
+    }
+    res.end();
   } else {
     res.statusCode = 404;
     res.end();
@@ -532,6 +547,20 @@ async function handleKinesisRequest(req, res, path) {
       res.statusCode = 500;
     }
     res.end();
+  } else if (path.includes('describe/my-stream')) {
+    try {
+      await kinesisClient.send(
+        new DescribeStreamCommand ({
+          StreamName: 'test_stream',
+          StreamARN: 'arn:aws:kinesis:us-west-2:000000000000:stream/test_stream' 
+        })
+      );
+      res.statusCode = 200;
+    } catch (err) {
+      console.log('Error describing the stream', err);
+      res.statusCode = 500;
+    }
+    res.end();
   } else {
     res.statusCode = 404;
     res.end();
@@ -1041,6 +1070,35 @@ async function handleLambdaRequest(req, res, path) {
   res.end();
 }
 
+async function handleCrossAccountRequest(req, res, path) {
+  const s3Client = new S3Client({
+    endpoint: _AWS_SDK_S3_ENDPOINT,
+    forcePathStyle: true,
+    region: "eu-central-1",
+    credentials: {
+      accessKeyId: "account_b_access_key_id",
+      secretAccessKey: "account_b_secret_access_key",
+      sessionToken: "account_b_token"
+    }
+  });
+
+  if (path.includes('createbucket/account_b')) {
+    try {
+      await s3Client.send(
+        new CreateBucketCommand({
+          Bucket: "cross-account-bucket",
+          CreateBucketConfiguration: { LocationConstraint: "eu-central-1" }
+        })
+      );
+      res.statusCode = 200;
+    } catch (err) {
+      console.log('Error creating a bucket in another account', err);
+      res.statusCode = 500;
+    }
+    res.end();
+  }
+}
+
 async function setupLambda() {
   const lambdaClient = new LambdaClient({
     endpoint: _AWS_SDK_ENDPOINT,

contract-tests/tests/test/amazon/aws-sdk/aws_sdk_test.py

@@ -19,7 +19,10 @@
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
-    AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER
+    AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
+    AWS_REMOTE_RESOURCE_ACCESS_KEY,
+    AWS_REMOTE_RESOURCE_ACCOUNT_ID,
+    AWS_REMOTE_RESOURCE_REGION
 )
 from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
 from opentelemetry.proto.metrics.v1.metrics_pb2 import ExponentialHistogramDataPoint, Metric
@@ -31,6 +34,7 @@
 
 _AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url"
 _AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name"
+_AWS_KINESIS_STREAM_ARN: str = "aws.kinesis.stream.arn"
 _AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name"
 _AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"
 _AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
@@ -49,6 +53,7 @@
 _GEN_AI_RESPONSE_FINISH_REASONS: str = "gen_ai.response.finish_reasons"
 _GEN_AI_USAGE_INPUT_TOKENS: str = 'gen_ai.usage.input_tokens'
 _GEN_AI_USAGE_OUTPUT_TOKENS: str = 'gen_ai.usage.output_tokens'
+_AWS_DYNAMODB_TABLE_ARN: str = "aws.dynamodb.table.arn"
 
 # pylint: disable=too-many-public-methods
 class AWSSDKTest(ContractTestBase):
@@ -236,6 +241,30 @@ def test_dynamodb_put_item(self):
             span_name="DynamoDB.PutItem",
         )
 
+    def test_dynamodb_describe_table(self):
+        self.do_test_requests(
+            "ddb/describe/some-table",
+            "GET",
+            200,
+            0,
+            0,
+            local_operation="GET /ddb",
+            remote_service="AWS::DynamoDB",
+            remote_operation="DescribeTable",
+            remote_resource_type="AWS::DynamoDB::Table",
+            remote_resource_identifier="put_test_table",
+            cloudformation_primary_identifier="put_test_table",
+            remote_resource_account_id="000000000000",
+            remote_resource_region="us-west-2",
+            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_error(self):
         self.do_test_requests(
             "ddb/error",
@@ -397,6 +426,28 @@ def test_kinesis_put_record(self):
             },
             span_name="Kinesis.PutRecord",
         )
+    
+    def test_kinesis_describe_stream(self):
+        self.do_test_requests(
+            "kinesis/describe/my-stream",
+            "GET",
+            200,
+            0,
+            0,
+            local_operation="GET /kinesis",
+            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",
+                _AWS_KINESIS_STREAM_ARN: "arn:aws:kinesis:us-west-2:000000000000:stream/test_stream"
+            },
+            span_name="Kinesis.DescribeStream",
+        )
 
     def test_kinesis_error(self):
         self.do_test_requests(
@@ -1034,6 +1085,27 @@ def test_lambda_error(self):
             span_name="Lambda.GetEventSourceMapping",
         )
 
+    def test_cross_account(self):
+        self.do_test_requests(
+            "cross-account/createbucket/account_b",
+            "GET",
+            200,
+            0,
+            0,
+            local_operation="GET /cross-account",
+            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: Need to add test_lambda_get_event_source_mapping once workaround is figured out for storing UUID between tests
 
     @override
@@ -1062,6 +1134,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_region", "None"),
+            kwargs.get("remote_resource_account_access_key", "None"),
         )
 
     def _assert_aws_attributes(
@@ -1073,7 +1148,10 @@ def _assert_aws_attributes(
         span_kind: str,
         remote_resource_type: str,
         remote_resource_identifier: str,
-        cloudformation_primary_identifier: str
+        cloudformation_primary_identifier: str,
+        remote_resource_account_id: str,
+        remote_resource_region: str,
+        remote_resource_account_access_key: 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())
@@ -1086,7 +1164,16 @@ def _assert_aws_attributes(
             self._assert_attribute(attributes_dict, AWS_REMOTE_RESOURCE_IDENTIFIER, remote_resource_identifier)
         if cloudformation_primary_identifier != "None":
             self._assert_attribute(attributes_dict, AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, cloudformation_primary_identifier)
-
+        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_account_access_key != "None":
+            assert remote_resource_identifier != "None"        
+            self._assert_str_attribute(attributes_dict, AWS_REMOTE_RESOURCE_ACCESS_KEY, remote_resource_account_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)
         self._assert_str_attribute(attributes_dict, AWS_SPAN_KIND, span_kind)
 
     @override
@@ -1187,10 +1274,23 @@ def _assert_metric_attributes(
             self._assert_attribute(attribute_dict, AWS_SPAN_KIND, kwargs.get("dependency_metric_span_kind") or "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_account_access_key = kwargs.get("remote_resource_account_access_key", "None")
+            remote_resource_region = kwargs.get("remote_resource_region", "None")
             if remote_resource_type != "None":
                 self._assert_attribute(attribute_dict, AWS_REMOTE_RESOURCE_TYPE, remote_resource_type)
             if remote_resource_identifier != "None":
                 self._assert_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_account_access_key != "None":
+                assert remote_resource_identifier != "None"        
+                self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_ACCESS_KEY, remote_resource_account_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)
@@ -1212,10 +1312,23 @@ def _assert_metric_attributes(
             self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, kwargs.get("dependency_metric_span_kind") or "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_account_access_key = kwargs.get("remote_resource_account_access_key", "None")
+            remote_resource_region = kwargs.get("remote_resource_region", "None")
             if remote_resource_type != "None":
                 self._assert_attribute(attribute_dict, AWS_REMOTE_RESOURCE_TYPE, remote_resource_type)
             if remote_resource_identifier != "None":
                 self._assert_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_account_access_key != "None":
+                assert remote_resource_identifier != "None"        
+                self._assert_str_attribute(attribute_dict, AWS_REMOTE_RESOURCE_ACCESS_KEY, remote_resource_account_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_service: Dict[str, AnyValue] = self._get_attributes_dict(service_dp.attributes)

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

@@ -19,3 +19,6 @@
 AWS_REMOTE_RESOURCE_IDENTIFIER: str = "aws.remote.resource.identifier"
 AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER: str = 'aws.remote.resource.cfn.primary.identifier'
 AWS_SPAN_KIND: str = "aws.span.kind"
+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"