query-orchestration/internal/query/queue/execute_test.go

package queryqueue_test

import (
	"context"
	"fmt"
	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/job/collector"
	queryprocessor "queryorchestration/internal/query/processor"
	queryqueue "queryorchestration/internal/query/queue"
	"queryorchestration/internal/query/result"
	"testing"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/pashagolub/pgxmock/v3"
	"github.com/stretchr/testify/assert"
)

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

	pool, err := pgxmock.NewPool()
	if err != nil {
		t.Fatalf("failed to open pgxmock database: %v", err)
	}
	queries := repository.New(pool)
	db := &database.Connection{
		Queries: queries,
		Pool:    pool,
	}
	coll := collector.Collector{
		ID: uuid.New(),
	}
	dbCollectorID := database.MustToDBUUID(coll.ID)

	queryOneID := uuid.New()
	queryOneVersion := int32(1)
	queryTwoID := uuid.New()
	queryTwoVersion := int32(2)
	queryThreeID := uuid.New()
	queryThreeVersion := int32(3)
	queryFourID := uuid.New()
	queryFourVersion := int32(4)
	queryFiveID := uuid.New()
	queryFiveVersion := int32(5)
	querySixID := uuid.New()
	querySixVersion := int32(6)
	contextID := uuid.New()
	contextVersion := int32(1)
	collectorQueries := []queryprocessor.Query{
		{ID: contextID, Type: queryprocessor.TypeContextFull, Version: contextVersion},
		{ID: queryOneID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{contextID}, Version: queryOneVersion},
		{ID: queryTwoID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{queryOneID}, Version: queryTwoVersion},
		{ID: queryThreeID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{queryOneID}, Version: queryThreeVersion},
		{ID: queryFourID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{contextID}, Version: queryFourVersion},
		{ID: queryFiveID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{querySixID}, Version: queryFiveVersion},
		{ID: querySixID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{contextID}, Version: querySixVersion},
	}

	rows := pgxmock.NewRows([]string{"collectorId", "queryId", "type", "queryVersion", "requiredIds"})
	for _, q := range collectorQueries {
		dbID := database.MustToDBUUID(q.ID)
		dbReqIDs := []pgtype.UUID{}
		if q.RequiredQueryIDs != nil {
			dbReqIDs = database.MustToDBUUIDArray(*q.RequiredQueryIDs)
		}
		ty, err := queryprocessor.ToDBQueryType(q.Type)
		assert.Nil(t, err)
		rows = rows.
			AddRow(dbCollectorID, dbID, ty, q.Version, dbReqIDs)
	}

	pool.ExpectQuery("name: ListCollectorQueries :many").WithArgs(dbCollectorID).WillReturnRows(rows)

	contextResultID := uuid.New()
	results := []*result.Result{
		{ID: contextResultID, QueryID: contextID, QueryVersion: contextVersion},
		{ID: uuid.New(), QueryID: queryFourID, QueryVersion: queryFourVersion},
		{ID: uuid.New(), QueryID: querySixID, QueryVersion: querySixVersion - 1},
		{ID: uuid.New(), QueryID: queryOneID, QueryVersion: queryOneVersion - 1},
	}

	expectedQueries := []*queryprocessor.Query{
		{ID: querySixID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{contextID}, Version: querySixVersion},
		{ID: queryFiveID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{querySixID}, Version: queryFiveVersion},
		{ID: queryOneID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{contextID}, Version: queryOneVersion},
		{ID: queryThreeID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{queryOneID}, Version: queryThreeVersion},
		{ID: queryTwoID, Type: queryprocessor.TypeJsonExtractor, RequiredQueryIDs: &[]uuid.UUID{queryOneID}, Version: queryTwoVersion},
	}

	docID := uuid.New()
	cleanVersion := int32(1)
	textVersion := int32(1)

	q, err := queryqueue.New(ctx, &queryqueue.NewConfig{
		db, &coll, results, docID, cleanVersion, textVersion,
	})
	assert.Nil(t, err)
	assert.Equal(t, expectedQueries, q.GetQueue())

	keyLayerOne := "key"
	keyLayerTwo := "key5"
	valueLayerTwo := "value"
	valueLayerOne := fmt.Sprintf("{\"%s\":\"%s\"}", keyLayerTwo, valueLayerTwo)
	valueContext := fmt.Sprintf("{\"%s\":%s}", keyLayerOne, valueLayerOne)

	pool.ExpectQuery("name: ListResultValuesByID :many").WithArgs([]pgtype.UUID{database.MustToDBUUID(contextResultID)}).WillReturnRows(
		pgxmock.NewRows([]string{"id", "queryId", "value"}).
			AddRow(database.MustToDBUUID(contextResultID), database.MustToDBUUID(contextID), valueContext),
	)
	pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(querySixID), querySixVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerOne))),
	)
	pool.ExpectQuery("name: SetResult :one").WithArgs(database.MustToDBUUID(querySixID), database.MustToDBUUID(docID), valueLayerOne, cleanVersion, textVersion, querySixVersion).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).AddRow(pgtype.UUID{}),
		)

	pool.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(
		pgxmock.NewRows([]string{"id", "queryId", "value"}).
			AddRow(pgtype.UUID{}, database.MustToDBUUID(querySixID), valueLayerOne),
	)
	pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(queryFiveID), queryFiveVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerTwo))),
	)
	pool.ExpectQuery("name: SetResult :one").WithArgs(database.MustToDBUUID(queryFiveID), database.MustToDBUUID(docID), valueLayerTwo, cleanVersion, textVersion, queryFiveVersion).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).AddRow(pgtype.UUID{}),
		)

	pool.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(
		pgxmock.NewRows([]string{"id", "queryId", "value"}).
			AddRow(database.MustToDBUUID(contextResultID), database.MustToDBUUID(contextID), valueContext),
	)
	pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(queryOneID), queryOneVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerOne))),
	)
	pool.ExpectQuery("name: SetResult :one").WithArgs(database.MustToDBUUID(queryOneID), database.MustToDBUUID(docID), valueLayerOne, cleanVersion, textVersion, queryOneVersion).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).AddRow(pgtype.UUID{}),
		)

	pool.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(
		pgxmock.NewRows([]string{"id", "queryId", "value"}).
			AddRow(pgtype.UUID{}, database.MustToDBUUID(queryOneID), valueLayerOne),
	)
	pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(queryThreeID), queryThreeVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerTwo))),
	)
	pool.ExpectQuery("name: SetResult :one").WithArgs(database.MustToDBUUID(queryThreeID), database.MustToDBUUID(docID), valueLayerTwo, cleanVersion, textVersion, queryThreeVersion).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).AddRow(pgtype.UUID{}),
		)

	pool.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(
		pgxmock.NewRows([]string{"id", "queryId", "value"}).
			AddRow(pgtype.UUID{}, database.MustToDBUUID(queryOneID), valueLayerOne),
	)
	pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(queryTwoID), queryTwoVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerTwo))),
	)
	pool.ExpectQuery("name: SetResult :one").WithArgs(database.MustToDBUUID(queryTwoID), database.MustToDBUUID(docID), valueLayerTwo, cleanVersion, textVersion, queryTwoVersion).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).AddRow(pgtype.UUID{}),
		)

	err = q.Execute(ctx)
	assert.Nil(t, err)
}