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

package queryqueue_test

import (
	"context"
	"errors"
	"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 TestService(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())
}

func TestQueueFail(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)

	dbErr := "database failure"
	pool.ExpectQuery("name: ListCollectorQueries :many").WithArgs(dbCollectorID).
		WillReturnError(errors.New(dbErr))

	results := []*result.Result{}

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

	_, err = queryqueue.New(ctx, &queryqueue.NewConfig{
		db, &coll, results, docID, cleanVersion, textVersion,
	})
	assert.EqualError(t, err, dbErr)
}