package doctextrunner_test import ( "fmt" "io" "regexp" "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" textractmock "queryorchestration/mocks/textract" "github.com/aws/aws-sdk-go-v2/service/s3" awstextract "github.com/aws/aws-sdk-go-v2/service/textract" "github.com/google/uuid" "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 TestDocTextRunner(t *testing.T) { cfg := &DocTextConfig{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint) cfg.QuerySyncURL = test.CreateQueue(t, cfg, test.QuerySyncRunnerName) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore mockTextract := textractmock.NewMockTextractClient(t) cfg.TextractClient = mockTextract runner := doctextrunner.New(&doctextrunner.Services{ Text: documenttext.New(cfg), }) err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: "clientid", Name: "client_name", }) require.NoError(t, err) docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{ Clientid: "clientid", Hash: "hash", }) require.NoError(t, err) part := uint16(1) filetype := "pdf" key := objectstore.BucketKey{ CreatedAt: time.Now().UTC(), ClientID: "clientid", EntityID: uuid.New(), Location: objectstore.Import, Part: &part, FileType: &filetype, } err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{ Documentid: docId, Bucket: "bucket", Key: key.String(), }) require.NoError(t, err) bucket := "bucket" hash := "hash" keyStr := key.String() cleanId, err := cfg.GetDBQueries().AddDocumentClean(t.Context(), &repository.AddDocumentCleanParams{ Documentid: docId, Bucket: &bucket, Hash: &hash, Key: &keyStr, Mimetype: repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.CleanmimetypeApplicationPdf, }, }) require.NoError(t, err) err = cfg.GetDBQueries().AddDocumentCleanEntry(t.Context(), &repository.AddDocumentCleanEntryParams{ Cleanid: cleanId, Version: 1, }) require.NoError(t, err) bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockStore.EXPECT(). GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { return *in.Bucket == "bucket" && *in.Key == key.String() }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: bodyMsg, }, nil) filename := "../../assets/textract/helloWorld.gen" textractBaseOut := test.GetTextractFileResponse(t, filename) mockTextract.EXPECT(). AnalyzeDocument( mock.Anything, mock.MatchedBy(func(in *awstextract.AnalyzeDocumentInput) bool { return true }), mock.Anything, ). Return(textractBaseOut[0]["base"], nil).Once() mockStore.EXPECT(). PutObject( mock.Anything, mock.MatchedBy(func(in *s3.PutObjectInput) bool { return *in.Bucket == "bucket" }), mock.Anything, ). Return(&s3.PutObjectOutput{}, nil) doc := doctextrunner.Body{ DocumentID: docId, } assert.True(t, runner.Process(t.Context(), doc)) test.AssertMessageBody(t, cfg, cfg.GetQuerySyncURL(), regexp.MustCompile(fmt.Sprintf("{\"id\":\"%s\"}", docId.String()))) } 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`