package documentclean import ( "context" "io" "strings" "testing" "queryorchestration/internal/database" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" 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.Document{ ID: uuid.New(), ClientID: uuid.New(), Hash: "example_hash", } inloc := document.Location{ Bucket: "bucket_name", Key: "/i/am/here", } dbid := database.MustToDBUUID(doc.ID) pool.ExpectQuery("name: GetDocument :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "hash"}). AddRow(dbid, database.MustToDBUUID(doc.ClientID), 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() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "mimetype", "fail"}), ) cleanid := database.MustToDBUUID(uuid.New()) 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, 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 == inloc.Bucket && *in.Key == inloc.Key }), 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 == inloc.Bucket && *in.Key == inloc.Key }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: body, }, nil) err = svc.clean(ctx, doc.ID) assert.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" inloc := document.Location{ Bucket: "bucket_name", Key: "/i/am/here", } mimeType := "application/pdf" mockS3.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) body := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockS3.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: body, }, nil) outloc, err := svc.executeCleanTasks(ctx, &CleanParams{ ID: id, Hash: hash, Location: inloc, }) assert.NoError(t, err) assert.EqualExportedValues(t, ExecuteCleanResponse{ location: &inloc, }, *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.Document{ ID: uuid.New(), } inloc := document.Location{ Bucket: "bucket_name", Key: "/i/am/here", } t.Run("no prev entry", func(t *testing.T) { mimeType := "application/pdf" params := &ExecuteCleanResponse{ location: &inloc, mimetype: (*MimeType)(&mimeType), } dbid := database.MustToDBUUID(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", "mimetype", "fail"}), ) cleanid := database.MustToDBUUID(uuid.New()) 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, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() err = svc.storeClean(ctx, doc.ID, params) assert.NoError(t, err) }) t.Run("diff prev entry", func(t *testing.T) { mimeType := "application/pdf" params := &ExecuteCleanResponse{ location: &inloc, mimetype: (*MimeType)(&mimeType), } dbid := database.MustToDBUUID(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", "mimetype", "fail"}). AddRow(database.MustToDBUUID(uuid.New()), dbid, nil, nil, int64(1), nil, repository.CleanfailtypeInvalidRead), ) cleanid := database.MustToDBUUID(uuid.New()) 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, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() err = svc.storeClean(ctx, doc.ID, params) assert.NoError(t, err) }) t.Run("same fail", func(t *testing.T) { reason := InvalidDocumentMimeType params := &ExecuteCleanResponse{ failReason: &reason, } dbid := database.MustToDBUUID(doc.ID) cleanid := database.MustToDBUUID(uuid.New()) pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "mimetype", "fail"}). AddRow(cleanid, dbid, nil, nil, int64(1), 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: database.MustToDBUUID(uuid.New()), Fail: repository.NullCleanfailtype{ Valid: true, Cleanfailtype: repository.CleanfailtypeInvalidMimetype, }, } reason := 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: database.MustToDBUUID(uuid.New()), Fail: repository.NullCleanfailtype{ Valid: true, Cleanfailtype: repository.CleanfailtypeInvalidMimetype, }, } reason := InvalidDocumentLargeDPI params := &ExecuteCleanResponse{ failReason: &reason, } assert.True(t, svc.isNewClean(lastEntry, params)) }) t.Run("same location", func(t *testing.T) { loc := document.Location{ Bucket: "bucket", Key: "hi", } mimeType := MimeTypePDF var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: database.MustToDBUUID(uuid.New()), Bucket: &loc.Bucket, Key: &loc.Key, Mimetype: repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), }, } params := &ExecuteCleanResponse{ location: &loc, mimetype: &mimeType, } assert.False(t, svc.isNewClean(lastEntry, params)) }) t.Run("different location", func(t *testing.T) { ogloc := document.Location{ Bucket: "bucket", Key: "hi", } newloc := document.Location{ Bucket: "bucket", Key: "different_location", } mimeType := MimeTypePDF var lastEntry *repository.GetMostRecentDocumentCleanEntryRow = &repository.GetMostRecentDocumentCleanEntryRow{ ID: database.MustToDBUUID(uuid.New()), Bucket: &ogloc.Bucket, Key: &ogloc.Key, Mimetype: repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), }, } params := &ExecuteCleanResponse{ location: &newloc, 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`