Cleaned up intent

* Added test to prove change * Reduced some boilerplate
open-telemetry · Mar 1, 2022 · 0c0e8b2 · 0c0e8b2
1 parent aa25924
commit 0c0e8b2
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,8 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 
-- `opentelemetry-instrumentation-aws-lambda` Added additional consumers. Fixes issue 902  
-  ([#926](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/926))
+- `opentelemetry-instrumentation-aws-lambda` `SpanKind.SERVER` by default, add more cases for `SpanKind.CONSUMER` services. ([#926](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/926))
 
 - `opentelemetry-instrumentation-dbapi` add experimental sql commenter capability
   ([#908](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/908))

diff --git a/...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py b/...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
@@ -187,17 +187,19 @@ def _instrumented_lambda_handler_call(
             lambda_event, event_context_extractor
         )
 
-        # See more:
-        # https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
-        # https://docs.aws.amazon.com/lambda/latest/dg/with-sns.html
-        # https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-content-structure.html
-        # https://docs.aws.amazon.com/lambda/latest/dg/with-ddb.html
-        span_kind = SpanKind.SERVER
-        valid_consumers = ["aws:sqs", "aws:s3", "aws:sns", "aws:dynamodb"]
-
+        span_kind = None
         try:
-            if lambda_event["Records"][0]["eventSource"] in valid_consumers:
+            if lambda_event["Records"][0]["eventSource"] in set(
+                ["aws:sqs", "aws:s3", "aws:sns", "aws:dynamodb"]
+            ):
+                # See more:
+                # https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
+                # https://docs.aws.amazon.com/lambda/latest/dg/with-sns.html
+                # https://docs.aws.amazon.com/AmazonS3/latest/userguide/notification-content-structure.html
+                # https://docs.aws.amazon.com/lambda/latest/dg/with-ddb.html
                 span_kind = SpanKind.CONSUMER
+            else:
+                span_kind = SpanKind.SERVER
         except (IndexError, KeyError, TypeError):
             span_kind = SpanKind.SERVER
 

diff --git a/.../opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py b/.../opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py
@@ -274,3 +274,29 @@ def test_lambda_no_error_with_invalid_flush_timeout(self):
         self.assertEqual(len(spans), 1)
 
         test_env_patch.stop()
+
+    def test_lambda_handles_multiple_consumers(self):
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                # NOT using the X-Ray Propagator
+                OTEL_PROPAGATORS: "tracecontext",
+            },
+        )
+        test_env_patch.start()
+
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda({"Records": [{"eventSource": "aws:sqs"}]})
+        mock_execute_lambda({"Records": [{"eventSource": "aws:s3"}]})
+        mock_execute_lambda({"Records": [{"eventSource": "aws:sns"}]})
+        mock_execute_lambda({"Records": [{"eventSource": "aws:dynamodb"}]})
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        test_env_patch.stop()