package queryrunner_test import ( "context" "fmt" "testing" queryrunner "queryorchestration/api/queryRunner" "queryorchestration/internal/database/repository" "queryorchestration/internal/query" "queryorchestration/internal/query/result" resultprocessor "queryorchestration/internal/query/result/processor" resultset "queryorchestration/internal/query/result/set" resultsync "queryorchestration/internal/query/result/sync" "queryorchestration/internal/server/runner" queryc "queryorchestration/internal/serviceconfig/queue/query" queuemock "queryorchestration/mocks/queue" "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) type QueryConfig struct { runner.BaseConfig[queryrunner.Body] queryc.QueryConfig } func TestQueryRunner(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, 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(&queryrunner.Services{ ResultSet: resultset.New(cfg, &resultset.Services{ Query: que, Result: result.New(cfg, &result.Services{ Query: que, }), Sync: resultsync.New(cfg), }), }) doc := queryrunner.Body{ DocumentID: uuid.New(), QueryID: uuid.New(), } qcfg := "{\"path\":\"examplekey\"}" reqQuery := uuid.New() query := &resultprocessor.Query{ ID: doc.QueryID, Version: 2, RequiredQueryIDs: &[]uuid.UUID{reqQuery}, Config: &qcfg, } params := &resultset.Set{ DocumentID: doc.DocumentID, } textEntryId := uuid.New() pool.ExpectQuery("name: GetTextEntryByDocId :one").WithArgs(params.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "hash", "cleanId", "extractionVersion"}). AddRow(textEntryId, params.DocumentID, "buket", "/i/am/here", "example", uuid.New(), int64(543)), ) pool.ExpectQuery("name: GetQuery :one").WithArgs(query.ID).WillReturnRows( pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}). AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs), ) pool.ExpectQuery("name: GetQueryWithVersion :one").WithArgs(&query.ID, &query.Version).WillReturnRows( pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}). AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs), ) requiredResultId := uuid.New() strVal := "{\"examplekey\":\"example_value\"}" pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, ¶ms.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"id", "queryId", "type", "value"}). AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal), ) pool.ExpectQuery("name: GetActiveQueryConfig :one").WithArgs(query.ID).WillReturnRows( pgxmock.NewRows([]string{"config"}). AddRow([]byte(qcfg)), ) pool.ExpectBegin() resultId := uuid.New() pool.ExpectQuery("name: AddResult :one").WithArgs(query.ID, pgxmock.AnyArg(), textEntryId, query.Version). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(resultId), ) pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, ¶ms.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"id", "queryId", "type", "value"}). AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal), ) pool.ExpectExec("name: AddResultDependency :exec").WithArgs(resultId, requiredResultId). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() pool.ExpectQuery("name: ListQueryDirectDependentsByDocumentID :many").WithArgs(&query.ID, &doc.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"queryId"}). AddRow(&reqQuery), ) 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) assert.True(t, runner.Process(ctx, doc)) }