package doccleanrunner_test import ( "context" "encoding/json" "fmt" "io" doccleanrunner "queryorchestration/api/docCleanRunner" "queryorchestration/internal/database" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" documentclean "queryorchestration/internal/document/clean" cleanversion "queryorchestration/internal/document/clean/version" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documenttext" objectstoremock "queryorchestration/mocks/objectstore" queuemock "queryorchestration/mocks/queue" "strings" "testing" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/sqs" "github.com/aws/aws-sdk-go-v2/service/sqs/types" "github.com/go-playground/validator/v10" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" ) type DocCleanConfig struct { serviceconfig.BaseConfig documenttext.DocTextConfig objectstore.ObjectStoreConfig } func TestDocCleanRunner(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocCleanConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore mockSQS := queuemock.NewMockSQSClient(t) cfg.QueueClient = mockSQS cfg.DocumentTextURL = "/i/am/here" runner := doccleanrunner.New(validator.New(), &doccleanrunner.Services{ Clean: documentclean.New(cfg, &documentclean.Services{ Version: cleanversion.New(cfg), }), }) assert.NotNil(t, runner) bod := doccleanrunner.Body{ ID: uuid.New(), } bodyBytes, err := json.Marshal(bod) assert.NoError(t, err) body := string(bodyBytes) msg := &types.Message{ Body: &body, } doc := document.Document{ ID: bod.ID, JobID: uuid.New(), Hash: "example_hash", } inloc := document.Location{ Bucket: "bucket_name", Key: "/i/am/here", } dbid := database.MustToDBUUID(doc.ID) pool.ExpectQuery("name: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"isclean"}). AddRow(false), ) pool.ExpectQuery("name: GetDocument :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "jobId", "hash"}). AddRow(dbid, database.MustToDBUUID(doc.JobID), doc.Hash), ) pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key"}). AddRow(dbid, inloc.Bucket, inloc.Key), ) mimeType := "application/pdf" dbmimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), } pool.ExpectBegin() cleanid := database.MustToDBUUID(uuid.New()) pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "mimetype", "fail"}), ) pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &inloc.Bucket, &inloc.Key, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, int32(1)). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.DocumentTextURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID.String()) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) mockStore.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ContentType: &mimeType, }, nil) bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockStore.EXPECT(). GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: bodyMsg, }, nil) err = runner.Process(ctx, msg) assert.NoError(t, err) } 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`