Files
query-orchestration/internal/document/clean/clean_test.go
T
Michael McGuinness 0815cb35fb Merged in feature/lint (pull request #87)
Basic Lint Checks

* basic
2025-03-05 12:05:46 +00:00

421 lines
11 KiB
Go

package documentclean
import (
"context"
"io"
"strings"
"testing"
"queryorchestration/internal/database"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document"
cleanversion "queryorchestration/internal/document/clean/version"
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,
svc: &Services{
Version: cleanversion.New(cfg),
},
}
doc := document.Document{
ID: uuid.New(),
JobID: 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", "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()
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, int32(1)).
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,
svc: &Services{
Version: cleanversion.New(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, int32(1)).
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, int32(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, int32(1)).
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, int32(1), nil, repository.NullCleanfailtype{
Valid: true,
Cleanfailtype: repository.CleanfailtypeInvalidMimetype,
}),
)
pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, int32(1)).
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`