query-orchestration/test/integration/queue/container_test.go

package queue_test

import (
	"context"
	"fmt"
	"strconv"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/sqs"
	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
	"github.com/docker/go-connections/nat"
	"github.com/stretchr/testify/assert"
	"github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/network"
	"github.com/testcontainers/testcontainers-go/wait"
)

type containerConfig struct {
	Queue   queueConfig
	DB      *dbConfig
	Network *testcontainers.DockerNetwork
}

func createContainer(t *testing.T, ctx context.Context, config containerConfig) testcontainers.Container {
	queueCredentials, err := config.Queue.Credentials.Retrieve(ctx)
	if err != nil {
		t.Fatal(err)
	}

	req := testcontainers.ContainerRequest{
		Image: "queryorchestration_queue:latest",
		Env: map[string]string{
			"QUEUE_URL":             config.Queue.URL,
			"AWS_DEFAULT_REGION":    config.Queue.Region,
			"AWS_REGION":            config.Queue.Region,
			"AWS_ACCESS_KEY_ID":     queueCredentials.AccessKeyID,
			"AWS_SECRET_ACCESS_KEY": queueCredentials.SecretAccessKey,
			"AWS_SESSION_TOKEN":     queueCredentials.SessionToken,
			"AWS_ENDPOINT_URL_SQS":  config.Queue.Endpoint,
			"DB_USER":               config.DB.User,
			"DB_PASS":               config.DB.Password,
			"DB_HOST":               config.DB.Host,
			"DB_NAME":               config.DB.Name,
			"DB_PORT":               strconv.Itoa(config.DB.Port),
		},
		Networks: []string{config.Network.Name},
	}

	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
	})
	if err != nil {
		t.Fatalf("Failed to start container: %v", err)
	}

	return container
}

type queueConfig struct {
	URL         string
	Region      string
	Endpoint    string
	Credentials aws.CredentialsProvider
}

type queue struct {
	Container *testcontainers.Container
	Client    *sqs.Client
	Config    *queueConfig
}

func createQueue(t *testing.T, ctx context.Context, network *testcontainers.DockerNetwork) *queue {
	queueName := "test-queue"
	region := "us-east-1"
	alias := "localstack"
	provider := credentials.NewStaticCredentialsProvider("test", "test", "")

	port, err := nat.NewPort("tcp", "4566")
	if err != nil {
		t.Fatalf("Failed to create port: %v", err)
	}

	req := testcontainers.ContainerRequest{
		Image: "localstack/localstack:latest",
		Env: map[string]string{
			"AWS_ACCESS_KEY_ID":      provider.Value.AccessKeyID,
			"AWS_SECRET_ACCESS_KEY":  provider.Value.SecretAccessKey,
			"AWS_SESSION_TOKEN":      provider.Value.SessionToken,
			"AWS_DEFAULT_REGION":     region,
			"AWS_REGION":             region,
			"SERVICES":               "sqs",
			"SKIP_SSL_CERT_DOWNLOAD": "1",
			"LOCALSTACK_HOST":        alias,
			"SQS_ENDPOINT_STRATEGY":  "path",
		},
		Networks: []string{network.Name},
		NetworkAliases: map[string][]string{
			network.Name: {alias},
		},
		ExposedPorts: []string{port.Port()},
		WaitingFor:   wait.ForListeningPort(port),
	}

	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
	})
	if err != nil {
		t.Fatalf("Failed to start container: %v", err)
	}

	host, err := container.Host(ctx)
	if err != nil {
		t.Fatalf("Failed to extract host: %v", err)
	}
	mappedPort, err := container.MappedPort(ctx, port)
	if err != nil {
		t.Fatalf("Failed to extract port: %v", err)
	}

	endpoint := fmt.Sprintf("http://%s:%s", host, mappedPort.Port())
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(provider), config.WithBaseEndpoint(endpoint))
	if err != nil {
		t.Fatal(err)
	}

	client := sqs.NewFromConfig(cfg)

	queueM, err := client.CreateQueue(ctx, &sqs.CreateQueueInput{
		QueueName: aws.String(queueName),
	})
	if err != nil {
		t.Fatal(err)
	}

	endpoint = fmt.Sprintf("http://%s:%s", alias, port.Port())
	queueC := &queueConfig{
		URL:         *queueM.QueueUrl,
		Region:      region,
		Credentials: provider,
		Endpoint:    endpoint,
	}

	return &queue{
		Container: &container,
		Client:    client,
		Config:    queueC,
	}
}

func createNetwork(t *testing.T, ctx context.Context) *testcontainers.DockerNetwork {
	network, err := network.New(ctx, network.WithDriver("bridge"))
	if err != nil {
		t.Fatal(err)
	}

	return network
}

type dbConfig struct {
	Name     string
	Host     string
	Port     int
	User     string
	Password string
}

func createDB(t *testing.T, ctx context.Context, network *testcontainers.DockerNetwork) (*dbConfig, testcontainers.Container) {
	alias := "postgres"
	port, err := nat.NewPort("tcp", "5432")
	if err != nil {
		t.Fatalf("Failed to create port: %v", err)
	}

	config := dbConfig{
		Name:     "queryorchestration",
		Password: "pass",
		User:     "postgres",
		Port:     port.Int(),
		Host:     alias,
	}

	req := testcontainers.ContainerRequest{
		Image: "postgres:latest",
		Env: map[string]string{
			"POSTGRES_DB":       config.Name,
			"POSTGRES_USER":     config.User,
			"POSTGRES_PASSWORD": config.Password,
		},
		Networks: []string{network.Name},
		NetworkAliases: map[string][]string{
			network.Name: {alias},
		},
		ExposedPorts: []string{port.Port()},
		WaitingFor:   wait.ForListeningPort(port),
	}

	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
		ContainerRequest: req,
		Started:          true,
	})
	if err != nil {
		t.Fatalf("Failed to start container: %v", err)
	}

	return &config, container
}

type message struct {
	URL  string
	Type string
	Body string
}

func sendMessage(t *testing.T, ctx context.Context, queue *sqs.Client, message message) {
	_, err := queue.SendMessage(ctx, &sqs.SendMessageInput{
		MessageAttributes: map[string]types.MessageAttributeValue{
			"type": {
				DataType:    aws.String("String"),
				StringValue: aws.String(message.Type),
			},
		},
		QueueUrl:    aws.String(message.URL),
		MessageBody: aws.String(message.Body),
	})
	if err != nil {
		t.Fatalf("Failed to send message: %v", err)
	}
}

func assertMessageWait(t *testing.T, ctx context.Context, queue *queue, msgType string) {
	time.Sleep(time.Second)

	result, err := queue.Client.ReceiveMessage(ctx, &sqs.ReceiveMessageInput{
		QueueUrl:            &queue.Config.URL,
		MaxNumberOfMessages: 1,
		WaitTimeSeconds:     5,
		VisibilityTimeout:   2,
		MessageAttributeNames: []string{
			"type",
		},
	})

	assert.Nil(t, err)
	for _, message := range result.Messages {
		assert.Equal(t, msgType, *message.MessageAttributes["type"].StringValue)
	}
}

func createDependencies(t *testing.T, ctx context.Context) (*queue, func()) {
	network := createNetwork(t, ctx)

	dbconfig, dbContainer := createDB(t, ctx, network)

	queue := createQueue(t, ctx, network)

	container := createContainer(t, ctx, containerConfig{
		Queue:   *queue.Config,
		DB:      dbconfig,
		Network: network,
	})

	return queue, func() {
		testcontainers.CleanupNetwork(t, network)
		dbContainer.Terminate(ctx)
		(*queue.Container).Terminate(ctx)
		container.Terminate(ctx)
	}
}