feat: update sns, add attributes and context params

aws/sns/sns.go

@@ -10,6 +10,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/sns"
+	"github.com/aws/aws-sdk-go-v2/service/sns/types"
 )
 
 type SNS struct {
@@ -66,12 +67,23 @@ func (s *SNS) Client() *sns.Client {
 
 // Publish publishes message to topicArn.
 func (s *SNS) Publish(topicArn string, message []byte) error {
+	return s.PublishWithContext(context.Background(), topicArn, message)
+}
+
+// PublishWithContext publishes message to topicArn using ctx.
+func (s *SNS) PublishWithContext(ctx context.Context, topicArn string, message []byte) error {
+	return s.PublishWithAttributes(ctx, topicArn, message, nil)
+}
+
+// PublishWithAttributes publishes message to topicArn with the supplied message attributes using ctx.
+func (s *SNS) PublishWithAttributes(ctx context.Context, topicArn string, message []byte, messageAttributes map[string]types.MessageAttributeValue) error {
 	input := sns.PublishInput{
-		Message:  aws.String(string(message)),
-		TopicArn: aws.String(topicArn),
+		Message:           aws.String(string(message)),
+		MessageAttributes: messageAttributes,
+		TopicArn:          aws.String(topicArn),
 	}
 
-	result, err := s.client.Publish(context.TODO(), &input)
+	result, err := s.client.Publish(ctx, &input)
 	if err != nil {
 		return err
 	}

aws/sns/sns_integration_test.go

@@ -4,13 +4,16 @@
 package sns
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
 	"os/exec"
 	"strings"
 	"testing"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sns/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -24,8 +27,12 @@ const (
 )
 
 type SNSMessage struct {
-	Message  string `json:"Message"`
-	TopicArn string `json:"TopicArn"`
+	Message           string `json:"Message"`
+	TopicArn          string `json:"TopicArn"`
+	MessageAttributes map[string]struct {
+		Type  string `json:"Type"`
+		Value string `json:"Value"`
+	} `json:"MessageAttributes"`
 }
 
 func setup(createQueue, createTopic bool) (string, string) {
@@ -167,6 +174,66 @@ func TestSNSPublish(t *testing.T) {
 	assert.Equal(t, testMessage, m.Message)
 }
 
+func TestSNSPublishWithContext(t *testing.T) {
+	// ARRANGE
+	queueArn, topicArn := setup(true, true)
+	defer teardown(queueArn, topicArn)
+	awsCmdSubscribeQueueToTopic(topicArn, queueArn)
+
+	client, err := New()
+	require.Nil(t, err, fmt.Sprintf("error creating sns client: %v", err))
+
+	// Normal publish with context
+	err = client.PublishWithContext(context.Background(), topicArn, []byte(testMessage))
+
+	// ASSERT
+	require.Nil(t, err)
+	var m SNSMessage
+	err = json.Unmarshal([]byte(awsCmdReceiveMessage(testQueue)), &m)
+	require.Nil(t, err)
+	assert.Equal(t, testMessage, m.Message)
+
+	// Cancelled context test
+	t.Run("cancelled", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		err := client.PublishWithContext(ctx, topicArn, []byte(testMessage))
+
+		require.Error(t, err)
+		assert.ErrorIs(t, err, context.Canceled)
+	})
+}
+
+func TestSNSPublishWithAttributes(t *testing.T) {
+	// ARRANGE
+	queueArn, topicArn := setup(true, true)
+	defer teardown(queueArn, topicArn)
+	awsCmdSubscribeQueueToTopic(topicArn, queueArn)
+
+	client, err := New()
+	require.Nil(t, err, fmt.Sprintf("error creating sns client: %v", err))
+
+	const attributeName = "test-attribute"
+	const attributeValue = "test-attribute-value"
+
+	// ACTION
+	err = client.PublishWithAttributes(context.Background(), topicArn, []byte(testMessage), map[string]types.MessageAttributeValue{
+		attributeName: {
+			DataType:    aws.String("String"),
+			StringValue: aws.String(attributeValue),
+		},
+	})
+
+	// ASSERT
+	require.Nil(t, err)
+	var m SNSMessage
+	err = json.Unmarshal([]byte(awsCmdReceiveMessage(testQueue)), &m)
+	require.Nil(t, err)
+	assert.Equal(t, testMessage, m.Message)
+	assert.Equal(t, attributeValue, m.MessageAttributes[attributeName].Value)
+}
+
 func TestSNSCreateTopic(t *testing.T) {
 	// ARRANGE
 	setup(false, false)