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

package runner

import (
	"context"
	"testing"
	"time"

	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue"
	"queryorchestration/internal/test"
	runnermock "queryorchestration/mocks/runner"

	"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) {
	cfg := &struct {
		BaseConfig[MockStruct]
		objectstore.ObjectStoreConfig
	}{}
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.QueueURL = test.CreateQueue(t, cfg, test.QueryRunnerName)

	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.Validator = validator.New()

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

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

func TestPollMessage(t *testing.T) {
	cfg := &struct {
		BaseConfig[MockStruct]
		objectstore.ObjectStoreConfig
	}{}
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.QueueURL = test.CreateQueue(t, cfg, test.QueryRunnerName)

	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.Validator = validator.New()

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

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

	err := cfg.SendToQueue(t.Context(), &queue.SendParams{
		QueueURL: cfg.GetQueueURL(),
		Body:     struct{}{},
	})
	require.NoError(t, err)

	err = ser.pollMessage(t.Context())
	require.NoError(t, err)

	ares, err := cfg.ReceiveFromQueue(t.Context(), &queue.ReceiveParams{
		QueueURL: cfg.GetQueueURL(),
	})
	require.NoError(t, err)
	require.Len(t, ares.Messages, 0)
}

func TestProcessMessage(t *testing.T) {
	cfg := &struct {
		BaseConfig[MockStruct]
		objectstore.ObjectStoreConfig
	}{}
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.QueueURL = test.CreateQueue(t, cfg, test.QueryRunnerName)

	mockController := runnermock.NewMockController[MockStruct](t)
	cfg.Controller = mockController
	cfg.Validator = validator.New()

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

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

	err := cfg.SendToQueue(t.Context(), &queue.SendParams{
		QueueURL: cfg.GetQueueURL(),
		Body:     struct{}{},
	})
	require.NoError(t, err)

	msg := test.AssertMessage(t, cfg, &queue.ReceiveParams{
		QueueURL: cfg.GetQueueURL(),
	})

	err = ser.processMessage(t.Context(), &msg)
	require.NoError(t, err)

	ares, err := cfg.ReceiveFromQueue(t.Context(), &queue.ReceiveParams{
		QueueURL: cfg.GetQueueURL(),
	})
	require.NoError(t, err)
	require.Len(t, ares.Messages, 0)
}

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))
	})
}