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

package document_test

import (
	"context"
	"gotemplate/internal/collector"
	"gotemplate/internal/database"
	"gotemplate/internal/database/repository"
	"gotemplate/internal/query"
	"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)
	collectorQueries := []query.Query{
		{ID: queryOneID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: uuid.Nil, Version: queryOneVersion},
		{ID: queryTwoID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryTwoVersion},
		{ID: queryThreeID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryThreeVersion},
		{ID: queryThreeID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: queryTwoID, Version: queryThreeVersion},
		{ID: queryFourID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: uuid.Nil, Version: queryFourVersion},
		{ID: queryFiveID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: querySixID, Version: queryFiveVersion},
		{ID: querySixID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: uuid.Nil, 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)
		rows = rows.
			AddRow(dbCollectorID, dbID, query.ToDBQueryType(q.Type), dbReqID, pgtype.Int4{Int32: int32(q.Version), Valid: true})
	}

	db.ExpectQuery("name: GetCollectorQueries :many").WithArgs(dbCollectorID).WillReturnRows(rows)

	results := []query.Result{
		{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.QueryTypeJsonExtractor, RequiredQueryID: uuid.Nil, Version: querySixVersion},
		{ID: queryFiveID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: querySixID, Version: queryFiveVersion},
		{ID: queryThreeID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: queryTwoID, Version: queryThreeVersion},
		{ID: queryTwoID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: queryOneID, Version: queryTwoVersion},
		{ID: queryOneID, Type: query.QueryTypeJsonExtractor, RequiredQueryID: uuid.Nil, Version: queryOneVersion},
	}

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

	q, err := query.NewQueue(ctx, queries, coll, &results, docID, cleanVersion, textVersion)
	assert.Nil(t, err)
	assert.Equal(t, expectedQueries, q.GetQueue())

	db.ExpectQuery("name: ListResultValuesByID :many").WithArgs([]pgtype.UUID{}).WillReturnRows(rows)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(querySixID), database.MustToDBUUID(docID), pgxmock.AnyArg(), cleanVersion, textVersion, querySixVersion).
		WillReturnResult(pgxmock.NewResult("", 1))
	db.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(rows)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(queryFiveID), database.MustToDBUUID(docID), pgxmock.AnyArg(), cleanVersion, textVersion, queryFiveVersion).
		WillReturnResult(pgxmock.NewResult("", 1))
	db.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(rows)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(queryThreeID), database.MustToDBUUID(docID), pgxmock.AnyArg(), cleanVersion, textVersion, queryThreeVersion).
		WillReturnResult(pgxmock.NewResult("", 1))
	db.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(rows)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(queryTwoID), database.MustToDBUUID(docID), pgxmock.AnyArg(), cleanVersion, textVersion, queryTwoVersion).
		WillReturnResult(pgxmock.NewResult("", 1))
	db.ExpectQuery("name: ListResultValuesByID :many").WithArgs(pgxmock.AnyArg()).WillReturnRows(rows)
	db.ExpectExec("name: SetResult :exec").WithArgs(pgxmock.AnyArg(), database.MustToDBUUID(queryOneID), database.MustToDBUUID(docID), pgxmock.AnyArg(), cleanVersion, textVersion, queryOneVersion).
		WillReturnResult(pgxmock.NewResult("", 1))

	err = q.Execute(ctx)
	assert.EqualError(t, err, "JSON Extraction requires 1 result")
}