Fix msg attribute size calcs, include data type attr (#103)

Co-authored-by: jwcorle <[email protected]>
co-go · Mar 5, 2024 · 5eae2d6 · 5eae2d6
1 parent 1caf917
commit 5eae2d6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/client.go b/client.go
@@ -192,6 +192,10 @@ func (c *Client) attributeSize(attributes map[string]types.MessageAttributeValue
 				sum.Add(int64(len(*attr.StringValue)))
 			}
 
+			if attr.DataType != nil {
+				sum.Add(int64(len(*attr.DataType)))
+			}
+
 			wg.Done()
 		}(k, v)
 	}

diff --git a/client_test.go b/client_test.go
@@ -142,13 +142,19 @@ func TestAttributeSize(t *testing.T) {
 		},
 	}))
 
+	assert.Equal(t, int64(20), c.attributeSize(map[string]types.MessageAttributeValue{
+		"testing_data_type": {
+			DataType: aws.String("int"),
+		},
+	}))
+
 	assert.Equal(t, int64(20), c.attributeSize(map[string]types.MessageAttributeValue{
 		"testing_binary": {
 			BinaryValue: []byte{1, 2, 3, 4, 5, 6},
 		},
 	}))
 
-	assert.Equal(t, int64(47), c.attributeSize(map[string]types.MessageAttributeValue{
+	assert.Equal(t, int64(65), c.attributeSize(map[string]types.MessageAttributeValue{
 		"binary_attr": {
 			BinaryValue: []byte{1, 2, 3, 4, 5, 6},
 		},
@@ -158,6 +164,9 @@ func TestAttributeSize(t *testing.T) {
 		"string_attr2": {
 			StringValue: aws.String("str"),
 		},
+		"data_type_attr1": {
+			DataType: aws.String("int"),
+		},
 	}))
 }