package documentclean import ( "context" "io" "strings" "testing" "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" documenttypes "queryorchestration/internal/document/types" "queryorchestration/internal/serviceconfig/objectstore" objectstoremock "queryorchestration/mocks/objectstore" "github.com/stretchr/testify/require" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" ) func TestClean(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &DocCleanConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockS3 := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockS3 svc := Service{ cfg: cfg, } doc := document.DocumentSummary{ ID: uuid.New(), ClientID: "yuyu", Hash: "example_hash", } bucket := "example_bucket" key := objectstore.BucketKey{ ClientID: doc.ClientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } keyStr := key.String() dbid := doc.ID pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "hash"}). AddRow(dbid, doc.ClientID, doc.Hash), ) pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key"}). AddRow(dbid, bucket, keyStr), ) mimeType := "application/pdf" dbmimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), } pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}), ) cleanid := uuid.New() pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &bucket, &keyStr, &doc.Hash, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() var length int64 = 10 mockS3.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == keyStr }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ContentType: &mimeType, ContentLength: &length, }, nil) body := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockS3.EXPECT(). GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { return *in.Bucket == bucket && *in.Key == keyStr }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: body, }, nil) err = svc.clean(ctx, doc.ID) require.NoError(t, err) } func TestExecuteCleanTasks(t *testing.T) { ctx := context.Background() cfg := &DocCleanConfig{} mockS3 := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockS3 svc := Service{ cfg: cfg, } id := uuid.New() hash := "example_hash" bucket := "example_bucket" key := objectstore.BucketKey{ ClientID: "client", CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } keyStr := key.String() mimeType := "application/pdf" mockS3.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == keyStr }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ContentType: &mimeType, }, nil) body := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockS3.EXPECT(). GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { return *in.Bucket == bucket && *in.Key == keyStr }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: body, }, nil) outloc, err := svc.executeCleanTasks(ctx, &CleanParams{ ID: id, Hash: hash, Bucket: bucket, Key: key, }) require.NoError(t, err) assert.EqualExportedValues(t, ExecuteCleanResponse{ hash: &hash, }, *outloc) } func TestStoreClean(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &DocCleanConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) svc := Service{ cfg: cfg, } doc := document.DocumentSummary{ ID: uuid.New(), } bucket := "example_bucket" key := objectstore.BucketKey{ ClientID: doc.ClientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } keyStr := key.String() t.Run("no prev entry", func(t *testing.T) { mimeType := "application/pdf" hash := "heyman" params := &ExecuteCleanResponse{ hash: &hash, mimetype: (*documenttypes.MimeType)(&mimeType), key: &key, bucket: &bucket, } dbid := doc.ID dbmimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), } pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}), ) cleanid := uuid.New() pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &bucket, &keyStr, params.hash, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() err = svc.storeClean(ctx, doc.ID, params) require.NoError(t, err) }) t.Run("diff prev entry", func(t *testing.T) { mimeType := "application/pdf" hash := "heyman" params := &ExecuteCleanResponse{ hash: &hash, bucket: &bucket, key: &key, mimetype: (*documenttypes.MimeType)(&mimeType), } dbid := doc.ID dbmimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), } pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}). AddRow(uuid.New(), dbid, nil, nil, int64(1), nil, nil, repository.CleanfailtypeInvalidRead), ) cleanid := uuid.New() pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &bucket, &keyStr, params.hash, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() err = svc.storeClean(ctx, doc.ID, params) require.NoError(t, err) }) t.Run("same fail", func(t *testing.T) { reason := documenttypes.InvalidDocumentMimeType params := &ExecuteCleanResponse{ failReason: &reason, } dbid := doc.ID cleanid := uuid.New() pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}). AddRow(cleanid, dbid, nil, nil, int64(1), nil, nil, repository.NullCleanfailtype{ Valid: true, Cleanfailtype: repository.CleanfailtypeInvalidMimetype, }), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() err = svc.storeClean(ctx, doc.ID, params) assert.EqualError(t, err, "invalid_mimetype") }) } func TestIsNewClean(t *testing.T) { svc := Service{} t.Run("no last entry", func(t *testing.T) { var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = nil params := &ExecuteCleanResponse{} assert.True(t, svc.isNewClean(lastEntry, params)) }) t.Run("same fail", func(t *testing.T) { var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: uuid.New(), Fail: repository.NullCleanfailtype{ Valid: true, Cleanfailtype: repository.CleanfailtypeInvalidMimetype, }, } reason := documenttypes.InvalidDocumentMimeType params := &ExecuteCleanResponse{ failReason: &reason, } assert.False(t, svc.isNewClean(lastEntry, params)) }) t.Run("different fail", func(t *testing.T) { var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: uuid.New(), Fail: repository.NullCleanfailtype{ Valid: true, Cleanfailtype: repository.CleanfailtypeInvalidMimetype, }, } reason := documenttypes.InvalidDocumentLargeDPI params := &ExecuteCleanResponse{ failReason: &reason, } assert.True(t, svc.isNewClean(lastEntry, params)) }) t.Run("same location", func(t *testing.T) { bucket := "example_bucket" key := objectstore.BucketKey{ ClientID: "diw", CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } keyStr := key.String() mimeType := documenttypes.MimeTypePDF var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: uuid.New(), Bucket: &bucket, Key: &keyStr, Mimetype: repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), }, } params := &ExecuteCleanResponse{ bucket: &bucket, key: &key, mimetype: &mimeType, } assert.False(t, svc.isNewClean(lastEntry, params)) }) t.Run("different location", func(t *testing.T) { bucket := "example_bucket" oldKey := objectstore.BucketKey{ ClientID: "hi", CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } oldKeyStr := oldKey.String() newKey := objectstore.BucketKey{ ClientID: "hi", CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } mimeType := documenttypes.MimeTypePDF var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: uuid.New(), Bucket: &bucket, Key: &oldKeyStr, Mimetype: repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), }, } params := &ExecuteCleanResponse{ bucket: &bucket, key: &newKey, mimetype: &mimeType, } assert.True(t, svc.isNewClean(lastEntry, params)) }) } 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`