Conti/add aws kinesis tests

manifests/cpp.yml

@@ -84,6 +84,8 @@ tests/:
     crossed_integrations/:
       test_kafka.py:
         Test_Kafka: missing_feature
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation: missing_feature
       test_sns_to_sqs.py:

manifests/dotnet.yml

@@ -209,6 +209,8 @@ tests/:
     crossed_integrations/:
       test_kafka.py:
         Test_Kafka: missing_feature
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation: missing_feature
       test_sns_to_sqs.py:

manifests/golang.yml

@@ -331,6 +331,10 @@ tests/:
         Test_Kafka:
           "*": irrelevant
           net-http: v0.1 # real version not known
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES:
+          "*": irrelevant
+          net-http: missing_feature (Endpoint not implemented)
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation:
           "*": irrelevant

manifests/java.yml

@@ -868,6 +868,10 @@ tests/:
         Test_Kafka:
           "*": irrelevant
           spring-boot: v0.1 # real version not known
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES:
+          "*": irrelevant
+          spring-boot: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation:
           "*": irrelevant

manifests/nodejs.yml

@@ -334,6 +334,10 @@ tests/:
         Test_Kafka:
           "*": irrelevant
           express4: v0.1 # real version not known
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES:
+          "*": irrelevant
+          express4: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation:
           "*": irrelevant

manifests/php.yml

@@ -179,6 +179,8 @@ tests/:
     crossed_integrations/:
       test_kafka.py:
         Test_Kafka: missing_feature
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation: missing_feature
       test_sns_to_sqs.py:

manifests/python.yml

@@ -501,6 +501,10 @@ tests/:
         Test_Kafka:
           '*': irrelevant
           flask-poc: v2.5.0
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES:
+          "*": irrelevant
+          flask-poc: v2.6.0
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation:
           '*': irrelevant

manifests/ruby.yml

@@ -218,6 +218,10 @@ tests/:
         Test_Kafka:
           "*": irrelevant
           rails70: v0.1 # real version unknown
+      test_kinesis.py:
+        Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES:
+          "*": irrelevant
+          rails70: missing_feature
       test_rabbitmq.py:
         Test_RabbitMQ_Trace_Context_Propagation:
           "*": irrelevant

tests/integrations/crossed_integrations/test_kinesis.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import json
+
+from tests.integrations.crossed_integrations.test_kafka import _python_buddy
+from utils import interfaces, scenarios, weblog, missing_feature, features
+from utils.tools import logger
+
+
+class _Test_Kinesis:
+    """Test Kinesis compatibility with inputted datadog tracer"""
+
+    BUDDY_TO_WEBLOG_STREAM = None
+    WEBLOG_TO_BUDDY_STREAM = None
+    buddy = None
+    buddy_interface = None
+
+    @classmethod
+    def get_span(cls, interface, span_kind, stream, operation):
+        logger.debug(f"Trying to find traces with span kind: {span_kind} and stream: {stream} in {interface}")
+
+        for data, trace in interface.get_traces():
+            for span in trace:
+                if not span.get("meta"):
+                    continue
+
+                if span["meta"].get("span.kind") not in span_kind:
+                    continue
+
+                # we want to skip all the kafka spans
+                if "aws.service" not in span["meta"] and "aws_service" not in span["meta"]:
+                    continue
+
+                if (
+                    "kinesis" not in span["meta"].get("aws.service", "").lower()
+                    and "kinesis" not in span["meta"].get("aws_service", "").lower()
+                ):
+                    continue
+
+                if operation.lower() != span["meta"].get("aws.operation", "").lower():
+                    continue
+
+                if operation.lower() == "getrecords" and span["meta"].get("language", "") == "javascript":
+                    # for nodejs we propagate from aws.response span which does not have the stream included on the span.
+                    if span["resource"] != "aws.response":
+                        continue
+                # elif stream != cls.get_stream_name(span):
+                #     continue
+
+                logger.debug(f"span found in {data['log_filename']}:\n{json.dumps(span, indent=2)}")
+                return span
+
+        logger.debug("No span found")
+        return None
+
+    @staticmethod
+    def get_stream(span) -> str | None:
+        """Extracts the stream from a span by trying various fields"""
+        stream = span["meta"].get("streamname", None)  # this is in nodejs, java, python
+
+        if stream is None:
+            if "aws.stream.url" in span["meta"]:
+                stream = span["meta"]["aws.stream.url"].split("/")[-1]
+            elif "messaging.url" in span["meta"]:
+                stream = span["meta"]["messaging.url"].split("/")[-1]
+
+            if stream is None:
+                logger.error(f"could not extract stream from this span:\n{span}")
+
+        return stream
+
+    def setup_produce(self):
+        """
+        send request A to weblog : this request will produce a Kinesis message
+        send request B to library buddy, this request will consume Kinesis message
+        """
+
+        self.production_response = weblog.get(
+            "/kinesis/produce", params={"stream": self.WEBLOG_TO_BUDDY_STREAM}, timeout=61
+        )
+        self.consume_response = self.buddy.get(
+            "/kinesis/consume", params={"stream": self.WEBLOG_TO_BUDDY_STREAM, "timeout": 60}, timeout=61
+        )
+
+    def test_produce(self):
+        """Check that a message produced to Kinesis is correctly ingested by a Datadog tracer"""
+
+        assert self.production_response.status_code == 200
+        assert self.consume_response.status_code == 200
+
+        # The weblog is the producer, the buddy is the consumer
+        self.validate_kinesis_spans(
+            producer_interface=interfaces.library,
+            consumer_interface=self.buddy_interface,
+            stream=self.WEBLOG_TO_BUDDY_STREAM,
+        )
+
+    @missing_feature(library="golang", reason="Expected to fail, Golang does not propagate context")
+    @missing_feature(library="ruby", reason="Expected to fail, Ruby does not propagate context")
+    @missing_feature(
+        library="java", reason="Expected to fail, Java defaults to using Xray headers to propagate context"
+    )
+    def test_produce_trace_equality(self):
+        """This test relies on the setup for produce, it currently cannot be run on its own"""
+        producer_span = self.get_span(
+            interfaces.library,
+            span_kind=["producer", "client"],
+            stream=self.WEBLOG_TO_BUDDY_STREAM,
+            operation="PutRecord",
+        )
+        consumer_span = self.get_span(
+            self.buddy_interface,
+            span_kind=["consumer", "client", "server"],
+            stream=self.WEBLOG_TO_BUDDY_STREAM,
+            operation="GetRecords",
+        )
+
+        # Both producer and consumer spans should be part of the same trace
+        # Different tracers can handle the exact propagation differently, so for now, this test avoids
+        # asserting on direct parent/child relationships
+        assert producer_span["trace_id"] == consumer_span["trace_id"]
+
+    def setup_consume(self):
+        """
+        send request A to library buddy : this request will produce a Kinesis message
+        send request B to weblog, this request will consume Kinesis message
+
+        request A: GET /library_buddy/produce_kinesis_message
+        request B: GET /weblog/consume_kinesis_message
+        """
+
+        self.production_response = self.buddy.get(
+            "/kinesis/produce", params={"stream": self.BUDDY_TO_WEBLOG_STREAM}, timeout=61
+        )
+        self.consume_response = weblog.get(
+            "/kinesis/consume", params={"stream": self.BUDDY_TO_WEBLOG_STREAM, "timeout": 60}, timeout=61
+        )
+
+    def test_consume(self):
+        """Check that a message by an app instrumented by a Datadog tracer is correctly ingested"""
+
+        assert self.production_response.status_code == 200
+        assert self.consume_response.status_code == 200
+
+        # The buddy is the producer, the weblog is the consumer
+        self.validate_kinesis_spans(
+            producer_interface=self.buddy_interface,
+            consumer_interface=interfaces.library,
+            stream=self.BUDDY_TO_WEBLOG_STREAM,
+        )
+
+    @missing_feature(library="golang", reason="Expected to fail, Golang does not propagate context")
+    @missing_feature(library="ruby", reason="Expected to fail, Ruby does not propagate context")
+    def test_consume_trace_equality(self):
+        """This test relies on the setup for consume, it currently cannot be run on its own"""
+        producer_span = self.get_span(
+            self.buddy_interface,
+            span_kind=["producer", "client"],
+            stream=self.BUDDY_TO_WEBLOG_STREAM,
+            operation="PutRecord",
+        )
+        consumer_span = self.get_span(
+            interfaces.library,
+            span_kind=["consumer", "client", "server"],
+            stream=self.BUDDY_TO_WEBLOG_STREAM,
+            operation="GetRecords",
+        )
+
+        # Both producer and consumer spans should be part of the same trace
+        # Different tracers can handle the exact propagation differently, so for now, this test avoids
+        # asserting on direct parent/child relationships
+        assert producer_span["trace_id"] == consumer_span["trace_id"]
+
+    def validate_kinesis_spans(self, producer_interface, consumer_interface, stream):
+        """
+        Validates production/consumption of Kinesis message.
+        It works the same for both test_produce and test_consume
+        """
+
+        producer_span = self.get_span(
+            producer_interface, span_kind=["producer", "client"], stream=stream, operation="PutRecord"
+        )
+        consumer_span = self.get_span(
+            consumer_interface, span_kind=["consumer", "client", "server"], stream=stream, operation="GetRecords"
+        )
+        # check that both consumer and producer spans exists
+        assert producer_span is not None
+        assert consumer_span is not None
+
+        # Assert that the consumer span is not the root
+        assert "parent_id" in consumer_span, "parent_id is missing in consumer span"
+
+        # returns both span for any custom check
+        return producer_span, consumer_span
+
+
+@scenarios.crossed_tracing_libraries
+@features.aws_kinesis_span_creationcontext_propagation_via_message_attributes_with_dd_trace
+class Test_Kinesis_PROPAGATION_VIA_MESSAGE_ATTRIBUTES(_Test_Kinesis):
+    buddy_interface = interfaces.python_buddy
+    buddy = _python_buddy
+    WEBLOG_TO_BUDDY_STREAM = "Test_Kinesis_propagation_via_message_attributes_weblog_to_buddy"
+    BUDDY_TO_WEBLOG_STREAM = "Test_Kinesis_propagation_via_message_attributes_buddy_to_weblog"

utils/_context/containers.py

@@ -752,10 +752,10 @@ def __init__(self, host_log_folder) -> None:
 class LocalstackContainer(TestedContainer):
     def __init__(self, host_log_folder) -> None:
         super().__init__(
-            image_name="localstack/localstack:3.0.2",
+            image_name="localstack/localstack:3.1.0",
             name="localstack-main",
             environment={
-                "LOCALSTACK_SERVICES": "sqs,sns,xray",
+                "LOCALSTACK_SERVICES": "kinesis,sqs,sns,xray",
                 "EXTRA_CORS_ALLOWED_HEADERS": "x-amz-request-id,x-amzn-requestid",
                 "EXTRA_CORS_EXPOSE_HEADERS": "x-amz-request-id,x-amzn-requestid",
                 "AWS_DEFAULT_REGION": "us-east-1",
@@ -764,6 +764,6 @@ def __init__(self, host_log_folder) -> None:
                 "DOCKER_HOST": "unix:///var/run/docker.sock",
             },
             host_log_folder=host_log_folder,
-            ports={4566: 4566},
+            ports={"4566": ("127.0.0.1", 4566)},
             volumes={"/var/run/docker.sock": {"bind": "/var/run/docker.sock", "mode": "rw"}},
         )

utils/_features.py

@@ -2006,3 +2006,13 @@ def host_user_managed_block_list(test_object):
         """
         pytest.mark.features(feature_id=276)(test_object)
         return test_object
+
+    @staticmethod
+    def aws_kinesis_span_creationcontext_propagation_via_message_attributes_with_dd_trace(test_object):
+        """
+        [AWS-Kinesis][Span Creation][Context Propagation] with dd-trace
+
+        https://feature-parity.us1.prod.dog/#/?feature=280
+        """
+        pytest.mark.features(feature_id=280)(test_object)
+        return test_object

utils/build/docker/python/flask-poc.Dockerfile

@@ -19,6 +19,8 @@ ENV _DD_APPSEC_DEDUPLICATION_ENABLED=false
 ENV DD_BOTOCORE_PROPAGATION_ENABLED=true
 ENV DD_KAFKA_PROPAGATION_ENABLED=true
 
+ENV LOG_LEVEL='DEBUG'
+
 # docker startup
 # FIXME: Ensure gevent patching occurs before ddtrace
 

utils/build/docker/python/flask/app.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import random
@@ -20,6 +21,8 @@
 from integrations.db.mssql import executeMssqlOperation
 from integrations.db.mysqldb import executeMysqlOperation
 from integrations.db.postgres import executePostgresOperation
+from integrations.messaging.aws.kinesis import kinesis_consume
+from integrations.messaging.aws.kinesis import kinesis_produce
 from integrations.messaging.aws.sns import sns_consume
 from integrations.messaging.aws.sns import sns_produce
 from integrations.messaging.aws.sqs import sqs_consume
@@ -259,6 +262,31 @@ def consume_sns_message():
         return output, 200
 
 
+@app.route("/kinesis/produce")
+def produce_kinesis_message():
+    stream = flask_request.args.get("stream", "DistributedTracing")
+    timeout = int(flask_request.args.get("timeout", 60))
+
+    # we only allow injection into JSON messages encoded as a string
+    message = json.dumps({"message": "Hello from Python Producer: Kinesis Context Propagation Test"})
+    output = kinesis_produce(stream, message, "1", timeout)
+    if "error" in output:
+        return output, 400
+    else:
+        return output, 200
+
+
+@app.route("/kinesis/consume")
+def consume_kinesis_message():
+    stream = flask_request.args.get("stream", "DistributedTracing")
+    timeout = int(flask_request.args.get("timeout", 60))
+    output = kinesis_consume(stream, timeout)
+    if "error" in output:
+        return output, 400
+    else:
+        return output, 200
+
+
 @app.route("/rabbitmq/produce")
 def produce_rabbitmq_message():
     queue = flask_request.args.get("queue", "DistributedTracingContextPropagation")