query-orchestration/internal/server/runner/poll_test.go

package runner

import (
	"context"
	"testing"
	"time"

	queuemock "queryorchestration/mocks/queue"
	runnermock "queryorchestration/mocks/runner"

	"github.com/aws/aws-sdk-go-v2/service/sqs"
	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
	"github.com/go-playground/validator/v10"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

type MockStruct struct{}

func TestPollMessages(t *testing.T) {
	ctx := context.Background()

	cfg := &BaseConfig[MockStruct]{}
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.QueueURL = "/i/am/here"
	cfg.Validator = validator.New()

	res := sqs.ReceiveMessageOutput{Messages: []types.Message{}}

	mockSQS.EXPECT().
		ReceiveMessage(
			mock.Anything,
			mock.MatchedBy(func(in *sqs.ReceiveMessageInput) bool {
				return *in.QueueUrl == cfg.QueueURL
			}),
			mock.Anything,
		).
		Return(&res, nil)

	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()

	scfg := &Server[MockStruct]{
		cfg:     cfg,
		cleanup: func() error { return nil },
	}
	scfg.Listen(ctx)
}

func TestPollMessage(t *testing.T) {
	ctx := context.Background()

	cfg := &BaseConfig[MockStruct]{}
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.QueueURL = "/i/am/here"
	cfg.Validator = validator.New()

	ser := &Server[MockStruct]{
		cfg:     cfg,
		cleanup: func() error { return nil },
	}

	id := "example_id"
	body := "{}"
	receipt := "receipts"

	mockSQS.EXPECT().
		ReceiveMessage(
			mock.Anything,
			mock.MatchedBy(func(in *sqs.ReceiveMessageInput) bool {
				return *in.QueueUrl == cfg.QueueURL
			}),
			mock.Anything,
		).
		Return(&sqs.ReceiveMessageOutput{
			Messages: []types.Message{
				{
					MessageId:     &id,
					Body:          &body,
					ReceiptHandle: &receipt,
				},
			},
		}, nil)

	mockController.EXPECT().
		Process(
			mock.Anything,
			mock.MatchedBy(func(in MockStruct) bool {
				return true
			}),
		).
		Return(true)

	mockSQS.EXPECT().
		DeleteMessage(
			mock.Anything,
			mock.MatchedBy(func(in *sqs.DeleteMessageInput) bool {
				return *in.QueueUrl == cfg.QueueURL && *in.ReceiptHandle == receipt
			}),
			mock.Anything,
		).
		Return(&sqs.DeleteMessageOutput{}, nil)

	err := ser.pollMessage(ctx)
	require.NoError(t, err)
}

func TestProcessMessage(t *testing.T) {
	ctx := context.Background()

	cfg := &BaseConfig[MockStruct]{}
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.QueueURL = "/i/am/here"
	cfg.Validator = validator.New()

	ser := &Server[MockStruct]{
		cfg:     cfg,
		cleanup: func() error { return nil },
	}

	id := "example_id"
	body := "{}"
	receipt := "receipts"
	msg := &types.Message{
		MessageId:     &id,
		Body:          &body,
		ReceiptHandle: &receipt,
	}

	mockController.EXPECT().
		Process(
			mock.Anything,
			mock.MatchedBy(func(in MockStruct) bool {
				return true
			}),
		).
		Return(true)

	mockSQS.EXPECT().
		DeleteMessage(
			mock.Anything,
			mock.MatchedBy(func(in *sqs.DeleteMessageInput) bool {
				return *in.QueueUrl == cfg.QueueURL && *in.ReceiptHandle == receipt
			}),
			mock.Anything,
		).
		Return(&sqs.DeleteMessageOutput{}, nil)

	err := ser.processMessage(ctx, msg)
	require.NoError(t, err)
}

func TestProcessBody(t *testing.T) {
	type MockStruct struct {
		Example string `json:"example" validate:"required"`
	}
	cfg := &BaseConfig[MockStruct]{}
	cfg.Validator = validator.New()

	ser := &Server[MockStruct]{
		cfg:     cfg,
		cleanup: func() error { return nil },
	}

	t.Run("valid body", func(t *testing.T) {
		id := "example_id"
		body := `{"example":"hi"}`
		msg := &types.Message{
			MessageId: &id,
			Body:      &body,
		}

		assert.NotNil(t, ser.processBody(msg))
	})
	t.Run("no msg", func(t *testing.T) {
		assert.Nil(t, ser.processBody(nil))
	})
	t.Run("no msg id", func(t *testing.T) {
		assert.Nil(t, ser.processBody(&types.Message{}))
	})
	t.Run("no body", func(t *testing.T) {
		id := "id"

		assert.Nil(t, ser.processBody(&types.Message{
			MessageId: &id,
		}))
	})
	t.Run("empty body", func(t *testing.T) {
		id := "example_id"
		body := ""
		msg := &types.Message{
			MessageId: &id,
			Body:      &body,
		}

		assert.Nil(t, ser.processBody(msg))
	})
	t.Run("invalid body", func(t *testing.T) {
		id := "example_id"
		body := "{}"
		msg := &types.Message{
			MessageId: &id,
			Body:      &body,
		}

		assert.Nil(t, ser.processBody(msg))
	})
}