package resultset import ( "context" "fmt" "testing" "queryorchestration/internal/database/repository" "queryorchestration/internal/query" "queryorchestration/internal/query/result" resultprocessor "queryorchestration/internal/query/result/processor" resultsync "queryorchestration/internal/query/result/sync" "queryorchestration/internal/serviceconfig" queryc "queryorchestration/internal/serviceconfig/queue/query" queuemock "queryorchestration/mocks/queue" "github.com/stretchr/testify/require" "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/mock" ) type ResultSetConfig struct { serviceconfig.BaseConfig queryc.QueryConfig } func TestSet(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &ResultSetConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockSQS := queuemock.NewMockSQSClient(t) cfg.QueueClient = mockSQS cfg.QueryURL = "/i/am/here" que := query.New(cfg) svc := Service{ cfg: cfg, svc: &Services{ Query: que, Result: result.New(cfg, &result.Services{ Query: que, }), Sync: resultsync.New(cfg), }, } qcfg := "{\"path\":\"examplekey\"}" requiredQuery := uuid.New() query := &resultprocessor.Query{ ID: uuid.New(), Version: 2, RequiredQueryIDs: &[]uuid.UUID{requiredQuery}, Config: &qcfg, } params := Set{ DocumentID: uuid.New(), QueryID: query.ID, } textEntryId := uuid.New() pool.ExpectQuery("name: GetTextEntryByDocId :one").WithArgs(params.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "hash", "cleanId", "triggerId", "textractJobId", "triggerVersion", "extractionVersion"}). AddRow(textEntryId, params.DocumentID, "buket", "/i/am/here", "example", uuid.New(), uuid.New(), "jobId", int64(123), 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(params.QueryID, params.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"queryId"}). AddRow(&requiredQuery), ) 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(), requiredQuery.String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err = svc.Set(ctx, ¶ms) require.NoError(t, err) } func TestInformQueryDependents(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &ResultSetConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockSQS := queuemock.NewMockSQSClient(t) cfg.QueueClient = mockSQS cfg.QueryURL = "/i/am/here" que := query.New(cfg) svc := Service{ cfg: cfg, svc: &Services{ Query: que, Result: result.New(cfg, &result.Services{ Query: que, }), Sync: resultsync.New(cfg), }, } reqOne := uuid.New() reqTwo := uuid.New() qIds := []uuid.UUID{ reqOne, reqTwo, } params := &Set{ DocumentID: uuid.New(), QueryID: uuid.New(), } pool.ExpectQuery("name: ListQueryDirectDependentsByDocumentID :many").WithArgs(params.QueryID, params.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"queryId"}). AddRow(&reqOne). AddRow(&reqTwo), ) 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(), qIds[0].String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) 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(), qIds[1].String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err = svc.informQueryDependents(ctx, params) require.NoError(t, err) }