package documenttext import ( "context" "fmt" "io" "os" "strings" "testing" "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/aws" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/querysync" "queryorchestration/internal/serviceconfig/textract" objectstoremock "queryorchestration/mocks/objectstore" queuemock "queryorchestration/mocks/queue" textractmock "queryorchestration/mocks/textract" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/sqs" awstextract "github.com/aws/aws-sdk-go-v2/service/textract" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) type DocTextConfig struct { aws.AWSConfig serviceconfig.BaseConfig querysync.QuerySyncConfig textract.TextractConfig objectstore.ObjectStoreConfig } func TestCreate(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &DocTextConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore mockTextract := textractmock.NewMockTextractClient(t) cfg.TextractClient = mockTextract mockQueue := queuemock.NewMockSQSClient(t) cfg.QueueClient = mockQueue cfg.QuerySyncURL = "ex" svc := Service{ cfg: cfg, } docId := uuid.New() t.Run("is not extracted", func(t *testing.T) { cleanId := uuid.New() bucket := "bucket" hash := `"a06d5e0e795adeb6d0b3732330dbcc15"` textId := uuid.New() key := objectstore.BucketKey{ ClientID: "aa", Location: objectstore.Import, CreatedAt: time.Now().UTC(), EntityID: cleanId, } keyStr := key.String() root, err := os.OpenRoot("../../../assets") require.NoError(t, err) textractBaseOut := getTextractFileResponse(t, "sampleGeneration/generated/helloWorld.gen", root) mimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.CleanmimetypeApplicationPdf, } pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(docId).WillReturnRows( pgxmock.NewRows([]string{"isextracted"}). AddRow(false), ) pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(docId). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}). AddRow(cleanId, docId, &bucket, &keyStr, int64(123), &hash, mimetype, nil, key.ClientID), ) bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockStore.EXPECT(). GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { return *in.Bucket == bucket && *in.Key == keyStr }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: bodyMsg, }, nil) sent := 0 mockTextract.EXPECT(). AnalyzeDocument( mock.Anything, mock.MatchedBy(func(in *awstextract.AnalyzeDocumentInput) bool { if sent == 0 { sent++ return true } return false }), mock.Anything, ). Return(&textractBaseOut["base"][0], nil).Once() pool.ExpectBegin() pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(textId), ) pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() mockQueue.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err = svc.Extract(ctx, docId) require.NoError(t, err) }) t.Run("is extracted", func(t *testing.T) { pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(docId).WillReturnRows( pgxmock.NewRows([]string{"isextracted"}). AddRow(true), ) mockQueue.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err = svc.Extract(ctx, docId) require.NoError(t, err) }) } func TestInformClean(t *testing.T) { ctx := context.Background() mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocTextConfig{} cfg.QueueClient = mockSQS cfg.QuerySyncURL = "/i/am/here" svc := Service{ cfg: cfg, } id := uuid.New() mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err := svc.informExtraction(ctx, id) require.NoError(t, err) }