package doctextrunner_test import ( "context" "fmt" "io" "strings" "testing" "time" doctextrunner "queryorchestration/api/docTextRunner" "queryorchestration/internal/database/repository" documenttext "queryorchestration/internal/document/text" "queryorchestration/internal/server/runner" "queryorchestration/internal/serviceconfig/aws" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/querysync" "queryorchestration/internal/serviceconfig/textract" "queryorchestration/internal/test" 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/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) type DocTextConfig struct { runner.BaseConfig[doctextrunner.Body] aws.AWSConfig querysync.QuerySyncConfig textract.TextractConfig objectstore.ObjectStoreConfig } func TestDocCleanRunner(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" runner := doctextrunner.New(&doctextrunner.Services{ Text: documenttext.New(cfg), }) assert.NotNil(t, runner) doc := doctextrunner.Body{ DocumentID: uuid.New(), } 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() filename := "../../assets/textract/helloWorld.gen" textractBaseOut := test.GetTextractFileResponse(t, filename) mimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.CleanmimetypeApplicationPdf, } pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(doc.DocumentID).WillReturnRows( pgxmock.NewRows([]string{"isextracted"}). AddRow(false), ) pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(doc.DocumentID). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}). AddRow(cleanId, doc.DocumentID, &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[0]["base"], 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\"}", doc.DocumentID.String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) assert.True(t, runner.Process(ctx, doc)) } const pdfHelloWorld = `%PDF-1.4 %���� 1 0 obj << /Type /Catalog /Pages 2 0 R /Version /1.4 >> endobj 2 0 obj << /Type /Pages /Kids [3 0 R] /Count 1 >> endobj 3 0 obj << /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 << /Type /Font /Subtype /Type1 /BaseFont /Helvetica >> >> >> /Contents 4 0 R >> endobj 4 0 obj << /Length 44 >> stream BT /F1 24 Tf 100 700 Td (Hello World!) Tj ET endstream endobj xref 0 5 0000000000 65535 f 0000000015 00000 n 0000000086 00000 n 0000000151 00000 n 0000000376 00000 n trailer << /Size 5 /Root 1 0 R >> startxref 472 %%EOF`