package queryrunner_test import ( "context" "encoding/json" "fmt" queryrunner "queryorchestration/api/queryRunner" "queryorchestration/internal/database" "queryorchestration/internal/database/repository" "queryorchestration/internal/query" "queryorchestration/internal/query/result" resultprocessor "queryorchestration/internal/query/result/processor" resultset "queryorchestration/internal/query/result/set" "queryorchestration/internal/server/runner" queryc "queryorchestration/internal/serviceconfig/queue/query" queuemock "queryorchestration/mocks/queue" "testing" "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/aws/aws-sdk-go-v2/service/sqs/types" "github.com/go-playground/validator/v10" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgtype" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" ) type QueryConfig struct { runner.BaseConfig queryc.QueryConfig } func TestQueryRunner(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &QueryConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockSQS := queuemock.NewMockSQSClient(t) cfg.QueueClient = mockSQS cfg.QueryURL = "/i/am/here" que := query.New(cfg) runner := queryrunner.New(validator.New(), &queryrunner.Services{ ResultSet: resultset.New(cfg, &resultset.Services{ Query: que, Result: result.New(cfg, &result.Services{ Query: que, }), }), }) assert.NotNil(t, runner) doc := resultset.Set{ DocumentID: uuid.New(), QueryID: uuid.New(), } bodyBytes, err := json.Marshal(doc) assert.NoError(t, err) body := string(bodyBytes) msg := &types.Message{ Body: &body, } qcfg := "{\"path\":\"examplekey\"}" query := &resultprocessor.Query{ ID: doc.QueryID, Version: 2, RequiredQueryIDs: &[]uuid.UUID{uuid.New()}, Config: &qcfg, } params := &resultset.Set{ DocumentID: doc.DocumentID, } pool.ExpectQuery("name: GetDocumentCleanEntry :one").WithArgs(database.MustToDBUUID(params.DocumentID)). WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key", "version"}). AddRow(database.MustToDBUUID(params.DocumentID), "buket", "/i/am/here", int32(1)), ) pool.ExpectQuery("name: GetDocumentTextEntry :one").WithArgs(database.MustToDBUUID(params.DocumentID)). WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key", "version"}). AddRow(database.MustToDBUUID(params.DocumentID), "buket", "/i/am/here", int32(1)), ) pool.ExpectQuery("name: GetQuery :one").WithArgs(database.MustToDBUUID(query.ID)).WillReturnRows( pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}). AddRow(database.MustToDBUUID(query.ID), repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), database.MustToDBUUIDArray(*query.RequiredQueryIDs)), ) pool.ExpectQuery("name: GetQueryWithVersion :one").WithArgs(database.MustToDBUUID(query.ID), query.Version).WillReturnRows( pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}). AddRow(database.MustToDBUUID(query.ID), repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), database.MustToDBUUIDArray(*query.RequiredQueryIDs)), ) pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(database.MustToDBUUID(params.DocumentID), database.MustToDBUUID(query.ID), query.Version). WillReturnRows( pgxmock.NewRows([]string{"queryId", "type", "value"}). AddRow(database.MustToDBUUID((*query.RequiredQueryIDs)[0]), repository.QuerytypeContextFull, "{\"examplekey\":\"example_value\"}"), ) pool.ExpectQuery("name: GetQueryConfig :one").WithArgs(database.MustToDBUUID(query.ID), query.Version).WillReturnRows( pgxmock.NewRows([]string{"id", "config"}). AddRow(pgtype.UUID{}, []byte(qcfg)), ) pool.ExpectExec("name: SetResult :exec").WithArgs(database.MustToDBUUID(query.ID), database.MustToDBUUID(params.DocumentID), pgxmock.AnyArg(), int32(1), int32(1), query.Version). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectQuery("name: ListQueryDirectDependentsByDocumentID :many").WithArgs(database.MustToDBUUID(query.ID), database.MustToDBUUID(doc.DocumentID)). WillReturnRows( pgxmock.NewRows([]string{"queryId"}). AddRow(database.MustToDBUUID((*query.RequiredQueryIDs)[0])), ) mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.QueryURL && *in.MessageBody == fmt.Sprintf("{\"document_id\":\"%s\",\"query_id\":\"%s\"}", params.DocumentID.String(), (*query.RequiredQueryIDs)[0].String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err = runner.Process(ctx, msg) assert.NoError(t, err) }