query-orchestration/test/unit/internal/queryQueue/queue_test.go

package document_test

import (
	"context"
	"fmt"
	"queryorchestration/internal/collector"
	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/query"
	"queryorchestration/internal/queryQueue"
	"queryorchestration/internal/result"
	"testing"

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

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

	db, err := pgxmock.NewConn()
	if err != nil {
		t.Fatalf("failed to open pgxmock database: %v", err)
	}
	defer db.Close(ctx)

	queries := repository.New(db)
	jobID := uuid.New()
	dbJobID := database.MustToDBUUID(jobID)
	collectorID := uuid.New()
	dbCollectorID := database.MustToDBUUID(collectorID)

	db.ExpectQuery("name: GetCollectorFromJobID :one").WithArgs(dbJobID).
		WillReturnRows(
			pgxmock.NewRows([]string{"id", "jobId", "minCleanVersion", "minTextVersion"}).
				AddRow(dbCollectorID, dbJobID, int32(1), int32(1)),
		)

	coll, err := collector.NewByJobId(ctx, queries, jobID)
	assert.Nil(t, err)

	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 := []query.Query{
		{ID: contextID, Type: query.TypeContextFull, RequiredQueryID: uuid.Nil, Version: contextVersion},
		{ID: queryOneID, Type: query.TypeJsonExtractor, RequiredQueryID: contextID, Version: queryOneVersion},
		{ID: queryTwoID, Type: query.TypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryTwoVersion},
		{ID: queryThreeID, Type: query.TypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryThreeVersion},
		{ID: queryFourID, Type: query.TypeJsonExtractor, RequiredQueryID: contextID, Version: queryFourVersion},
		{ID: queryFiveID, Type: query.TypeJsonExtractor, RequiredQueryID: querySixID, Version: queryFiveVersion},
		{ID: querySixID, Type: query.TypeJsonExtractor, RequiredQueryID: contextID, Version: querySixVersion},
	}

	rows := pgxmock.NewRows([]string{"collectorId", "queryId", "type", "requiredQueryId", "queryVersion"})
	for _, q := range collectorQueries {
		dbID := database.MustToDBUUID(q.ID)
		dbReqID := database.MustToDBUUID(q.RequiredQueryID)
		ty, err := query.ToDBQueryType(q.Type)
		assert.Nil(t, err)
		rows = rows.
			AddRow(dbCollectorID, dbID, ty, dbReqID, pgtype.Int4{Int32: int32(q.Version), Valid: true})
	}

	db.ExpectQuery("name: GetCollectorQueries :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 := []query.Query{
		{ID: querySixID, Type: query.TypeJsonExtractor, RequiredQueryID: contextID, Version: querySixVersion},
		{ID: queryFiveID, Type: query.TypeJsonExtractor, RequiredQueryID: querySixID, Version: queryFiveVersion},
		{ID: queryOneID, Type: query.TypeJsonExtractor, RequiredQueryID: contextID, Version: queryOneVersion},
		{ID: queryThreeID, Type: query.TypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryThreeVersion},
		{ID: queryTwoID, Type: query.TypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryTwoVersion},
	}

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

	q, err := queryQueue.New(ctx, queries, 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)

	db.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),
	)
	db.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(querySixID), querySixVersion).WillReturnRows(
		pgxmock.NewRows([]string{"id", "config"}).
			AddRow(pgtype.UUID{}, []byte(fmt.Sprintf("{\"path\":\"%s\"}", keyLayerOne))),
	)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(querySixID), database.MustToDBUUID(docID), valueLayerOne, cleanVersion, textVersion, querySixVersion).
		WillReturnResult(pgxmock.NewResult("", 1))

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

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

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

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

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