From 033cc1f7ed09c33e401b9514ed30d1160cf58899 Mon Sep 17 00:00:00 2001
From: Yaniv Davidi <yaniv@aspecto.io>
Date: Fri, 14 Jan 2022 14:57:30 +0200
Subject: [PATCH] fix(aws-sdk): calc propagation fields count before context
 inject (#738)

* chore: calc propagation fields count before inject

* chore(aws-sdk): lint:fix

* chore(aws-sdk): cover MessageAttributes with tests

Co-authored-by: Amir Blum <amir@aspecto.io>
---
 .../src/services/MessageAttributes.ts         |  5 +-
 .../test/MessageAttributes.test.ts            | 79 +++++++++++++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 plugins/node/opentelemetry-instrumentation-aws-sdk/test/MessageAttributes.test.ts

diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/MessageAttributes.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/MessageAttributes.ts
index b077ba9eb9..45aa2b617f 100644
--- a/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/MessageAttributes.ts
+++ b/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/MessageAttributes.ts
@@ -71,7 +71,10 @@ export const injectPropagationContext = (
   attributesMap?: SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap
 ): SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap => {
   const attributes = attributesMap ?? {};
-  if (Object.keys(attributes).length < MAX_MESSAGE_ATTRIBUTES) {
+  if (
+    Object.keys(attributes).length + propagation.fields().length <=
+    MAX_MESSAGE_ATTRIBUTES
+  ) {
     propagation.inject(context.active(), attributes, contextSetter);
   } else {
     diag.warn(
diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/MessageAttributes.test.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/MessageAttributes.test.ts
new file mode 100644
index 0000000000..9cb8cfdde5
--- /dev/null
+++ b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/MessageAttributes.test.ts
@@ -0,0 +1,79 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as expect from 'expect';
+import {
+  MAX_MESSAGE_ATTRIBUTES,
+  contextSetter,
+  injectPropagationContext,
+} from '../src/services/MessageAttributes';
+
+describe('MessageAttributes', () => {
+  describe('MAX_MESSAGE_ATTRIBUTES', () => {
+    it('should be 10', () => {
+      expect(MAX_MESSAGE_ATTRIBUTES).toBe(10);
+    });
+  });
+
+  describe('contextSetter', () => {
+    it('should set parent context in sqs receive callback', () => {
+      const contextKey = 'key';
+      const contextValue = 'value';
+      const contextCarrier = {};
+      contextSetter.set(contextCarrier, contextKey, contextValue);
+
+      const expectedContext = {
+        [contextKey]: { DataType: 'String', StringValue: contextValue },
+      };
+      expect(contextCarrier).toStrictEqual(expectedContext);
+    });
+  });
+
+  describe('injectPropagationContext', () => {
+    it('should inject context if there are available attributes', () => {
+      const contextAttributes = {
+        key1: { DataType: 'String', StringValue: 'value1' },
+        key2: { DataType: 'String', StringValue: 'value2' },
+        key3: { DataType: 'String', StringValue: 'value3' },
+        key4: { DataType: 'String', StringValue: 'value4' },
+        key5: { DataType: 'String', StringValue: 'value5' },
+      };
+
+      expect(Object.keys(contextAttributes).length).toBe(5);
+      injectPropagationContext(contextAttributes);
+      expect(Object.keys(contextAttributes).length).toBeGreaterThan(5);
+    });
+
+    it('should not inject context if there not enough available attributes', () => {
+      const contextAttributes = {
+        key1: { DataType: 'String', StringValue: 'value1' },
+        key2: { DataType: 'String', StringValue: 'value2' },
+        key3: { DataType: 'String', StringValue: 'value3' },
+        key4: { DataType: 'String', StringValue: 'value4' },
+        key5: { DataType: 'String', StringValue: 'value5' },
+        key6: { DataType: 'String', StringValue: 'value6' },
+        key7: { DataType: 'String', StringValue: 'value7' },
+        key8: { DataType: 'String', StringValue: 'value8' },
+        key9: { DataType: 'String', StringValue: 'value9' },
+        key10: { DataType: 'String', StringValue: 'value10' },
+      };
+
+      expect(Object.keys(contextAttributes).length).toBe(10);
+      injectPropagationContext(contextAttributes);
+      expect(Object.keys(contextAttributes).length).toBe(10);
+    });
+  });
+});