query-orchestration/internal/test/queue.go

package test

import (
	"context"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/queue"
	"regexp"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/sqs"
	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
	"github.com/stretchr/testify/assert"
)

func CreateQueue(t *testing.T, ctx context.Context, cfg serviceconfig.ConfigProvider, name string) string {
	queueM, err := cfg.GetQueueClient().CreateQueue(ctx, &sqs.CreateQueueInput{
		QueueName: aws.String(name),
	})
	if err != nil {
		t.Fatal(err)
	}

	err = cfg.PingQueueByURL(ctx, *queueM.QueueUrl)
	if err != nil {
		t.Fatal(err)
	}

	return *queueM.QueueUrl
}

func AssertMessage(t *testing.T, ctx context.Context, cfg serviceconfig.ConfigProvider, params *queue.ReceiveParams) types.Message {
	result, err := cfg.ReceiveFromQueue(ctx, params)
	assert.NoError(t, err)
	assert.NotNil(t, result.Messages)
	assert.Len(t, result.Messages, 1)
	assert.NotNil(t, result.Messages[0])

	return result.Messages[0]
}

func AssertMessageBody(t *testing.T, ctx context.Context, cfg serviceconfig.ConfigProvider, url string, body *regexp.Regexp) {
	message := AssertMessage(t, ctx, cfg, &queue.ReceiveParams{
		QueueURL: url,
	})

	assert.Regexp(t, body, *message.Body)
}

func AssertMessageAttr(t *testing.T, ctx context.Context, cfg serviceconfig.ConfigProvider, url string, name string, value *regexp.Regexp) {
	message := AssertMessage(t, ctx, cfg, &queue.ReceiveParams{
		QueueURL:   url,
		Attributes: []string{name},
	})

	assert.NotNil(t, message.MessageAttributes[name])
	assert.Regexp(t, value, *(message.MessageAttributes[name]).StringValue)
}