Merged in feature/remove-mocks (pull request #202)
Feature/remove mocks * remove mocks * cleanup db migrations
This commit is contained in:
Jay Brown
@@ -1,445 +0,0 @@
|
||||
package documentclean
|
||||
|
||||
import (
|
||||
"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 := t.Context()
|
||||
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 := t.Context()
|
||||
|
||||
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 := t.Context()
|
||||
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`
|
||||
@@ -1,493 +0,0 @@
|
||||
package documentclean
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
documenttypes "queryorchestration/internal/document/types"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetBodyMimeType(t *testing.T) {
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := "invalid"
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("pdf", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := fmt.Sprintf("%s %s", PDFSignatureStr, PDFSignatureEndStr)
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypePDF, contentType)
|
||||
})
|
||||
t.Run("pdf start", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := fmt.Sprintf("%s---------------", PDFSignatureStr)
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("pdf end", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := fmt.Sprintf("%s---------------", PDFSignatureEndStr)
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("short", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := ""
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
_, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
|
||||
assert.Error(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetMetadataMimeType(t *testing.T) {
|
||||
t.Run("nil type", func(t *testing.T) {
|
||||
svc := Service{}
|
||||
|
||||
contentType := svc.getMetadataMimeType(nil)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
svc := Service{}
|
||||
|
||||
inType := "invalid_type"
|
||||
contentType := svc.getMetadataMimeType(&inType)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("pdf", func(t *testing.T) {
|
||||
svc := Service{}
|
||||
|
||||
inType := "application/pdf"
|
||||
contentType := svc.getMetadataMimeType(&inType)
|
||||
assert.Equal(t, documenttypes.MimeTypePDF, contentType)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetAcceptedMimeType(t *testing.T) {
|
||||
t.Run("invalid type", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := "invalid"
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String()
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
inType := "invalid_type"
|
||||
contentType, err := svc.getAcceptedMimeType(ctx, params, &inType, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypeInvalid, contentType)
|
||||
})
|
||||
t.Run("pdf mimetype", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
inType := "application/pdf"
|
||||
contentType, err := svc.getAcceptedMimeType(ctx, &CleanParams{}, &inType, 0)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypePDF, contentType)
|
||||
})
|
||||
t.Run("pdf format", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
|
||||
bodyStr := fmt.Sprintf("%s %s", PDFSignatureStr, PDFSignatureEndStr)
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String()
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
inType := "invalid_type"
|
||||
contentType, err := svc.getAcceptedMimeType(ctx, params, &inType, int64(len(bodyStr)))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, documenttypes.MimeTypePDF, contentType)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetBytesBuffer(t *testing.T) {
|
||||
t.Run("valid", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
start := int64(0)
|
||||
bodyStr := "invalid"
|
||||
length := int64(len(bodyStr))
|
||||
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash && *in.Range == "bytes=0-6"
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, buffer)
|
||||
})
|
||||
t.Run("negative start", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
start := int64(-1)
|
||||
bodyStr := "invalid"
|
||||
length := int64(len(bodyStr))
|
||||
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash && *in.Range == "bytes=0-6"
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
|
||||
require.NoError(t, err)
|
||||
assert.NotNil(t, buffer)
|
||||
})
|
||||
t.Run("zero length", func(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := &DocCleanConfig{}
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
params := &CleanParams{
|
||||
Hash: "example_hash",
|
||||
Bucket: "example_bucket",
|
||||
Key: objectstore.BucketKey{
|
||||
ClientID: "ddddddd",
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
},
|
||||
}
|
||||
start := int64(0)
|
||||
bodyStr := ""
|
||||
length := int64(0)
|
||||
|
||||
body := io.NopCloser(strings.NewReader(bodyStr))
|
||||
mockS3.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == params.Bucket && *in.Key == params.Key.String() && *in.IfMatch == params.Hash && *in.Range == "bytes=0-0"
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, buffer)
|
||||
})
|
||||
}
|
||||
@@ -1,15 +1,17 @@
|
||||
// Package documentclean provides document cleaning functionality.
|
||||
// Note: Text extraction pipeline has been removed. The document processing
|
||||
// pipeline now ends after cleaning. See remove_texttract_and_mocks_plan.md.
|
||||
package documentclean
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
doctextrunner "queryorchestration/api/docTextRunner"
|
||||
"queryorchestration/internal/serviceconfig/queue"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Clean processes a document by cleaning it. Previously this also triggered
|
||||
// text extraction, but that feature has been removed.
|
||||
func (s *Service) Clean(ctx context.Context, documentId uuid.UUID) error {
|
||||
isclean, err := s.cfg.GetDBQueries().HasDocumentCleanEntry(ctx, documentId)
|
||||
if err != nil {
|
||||
@@ -23,24 +25,6 @@ func (s *Service) Clean(ctx context.Context, documentId uuid.UUID) error {
|
||||
}
|
||||
}
|
||||
|
||||
err = s.informClean(ctx, documentId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) informClean(ctx context.Context, documentId uuid.UUID) error {
|
||||
err := s.cfg.SendToQueue(ctx, &queue.SendParams{
|
||||
QueueURL: s.cfg.GetDocumentTextURL(),
|
||||
Body: doctextrunner.Body{
|
||||
DocumentID: documentId,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Note: Text extraction has been removed. Document pipeline ends here.
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1,164 +0,0 @@
|
||||
package documentclean
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/document"
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/aws"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/queue/documenttext"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
queuemock "queryorchestration/mocks/queue"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/aws-sdk-go-v2/service/sqs"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pashagolub/pgxmock/v3"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type DocCleanConfig struct {
|
||||
aws.AWSConfig
|
||||
serviceconfig.BaseConfig
|
||||
documenttext.DocTextConfig
|
||||
objectstore.ObjectStoreConfig
|
||||
}
|
||||
|
||||
func TestCreate(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
|
||||
mockSQS := queuemock.NewMockSQSClient(t)
|
||||
|
||||
cfg := &DocCleanConfig{}
|
||||
cfg.QueueClient = mockSQS
|
||||
cfg.DocumentTextURL = "/i/am/here"
|
||||
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: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"isclean"}).
|
||||
AddRow(false),
|
||||
)
|
||||
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()
|
||||
|
||||
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)
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&sqs.SendMessageOutput{}, nil)
|
||||
|
||||
mockS3.EXPECT().
|
||||
HeadObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == key.String()
|
||||
}),
|
||||
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 == key.String()
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: body,
|
||||
}, nil)
|
||||
|
||||
err = svc.Clean(ctx, doc.ID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestInformClean(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
mockSQS := queuemock.NewMockSQSClient(t)
|
||||
|
||||
cfg := &DocCleanConfig{}
|
||||
cfg.QueueClient = mockSQS
|
||||
cfg.DocumentTextURL = "/i/am/here"
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
id := uuid.New()
|
||||
|
||||
mockSQS.EXPECT().
|
||||
SendMessage(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
|
||||
return *in.QueueUrl == cfg.DocumentTextURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id)
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&sqs.SendMessageOutput{}, nil)
|
||||
|
||||
err := svc.informClean(ctx, id)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
@@ -1,21 +1,26 @@
|
||||
// Package documentclean provides document cleaning functionality.
|
||||
// Note: Text extraction pipeline has been removed. The document processing
|
||||
// pipeline now ends after cleaning.
|
||||
package documentclean
|
||||
|
||||
import (
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/queue/documenttext"
|
||||
)
|
||||
|
||||
// ConfigProvider defines the configuration requirements for the document clean service.
|
||||
// Note: documenttext.ConfigProvider has been removed since text extraction is deprecated.
|
||||
type ConfigProvider interface {
|
||||
serviceconfig.ConfigProvider
|
||||
documenttext.ConfigProvider
|
||||
objectstore.ConfigProvider
|
||||
}
|
||||
|
||||
// Service provides document cleaning operations.
|
||||
type Service struct {
|
||||
cfg ConfigProvider
|
||||
}
|
||||
|
||||
// New creates a new document cleaning service.
|
||||
func New(cfg ConfigProvider) *Service {
|
||||
return &Service{
|
||||
cfg: cfg,
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
// Package documentclean_test contains tests for the document clean service.
|
||||
// Note: Text extraction has been removed. The document pipeline ends after cleaning.
|
||||
package documentclean_test
|
||||
|
||||
import (
|
||||
@@ -6,14 +8,12 @@ import (
|
||||
documentclean "queryorchestration/internal/document/clean"
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/queue/documenttext"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
type DocCleanConfig struct {
|
||||
serviceconfig.BaseConfig
|
||||
documenttext.DocTextConfig
|
||||
objectstore.ConfigProvider
|
||||
}
|
||||
|
||||
|
||||
@@ -1,449 +1,21 @@
|
||||
// Package documentinit tests for document initialization.
|
||||
// Note: Mock-based tests have been removed. See remove_texttract_and_mocks_plan.md.
|
||||
package documentinit
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/document"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/test"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
queuemock "queryorchestration/mocks/queue"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/aws-sdk-go-v2/service/sqs"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pashagolub/pgxmock/v3"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCreate(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
mockSQS := queuemock.NewMockSQSClient(t)
|
||||
cfg.QueueClient = mockSQS
|
||||
|
||||
// Set up mock S3 client for file size measurement
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
cfg.DocumentSyncURL = "/i/am/here"
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
clientId := "hi"
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: clientId,
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
// Mock HeadObject to return file size
|
||||
fileSize := int64(1024)
|
||||
mockS3.EXPECT().
|
||||
HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == key.String()
|
||||
})).
|
||||
Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
|
||||
|
||||
pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}),
|
||||
)
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(doc.ID),
|
||||
)
|
||||
pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(doc.ID, bucket, key.String()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(doc.ID).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "clientId", "hash"}).
|
||||
AddRow(doc.ID, doc.ClientID, doc.Hash),
|
||||
)
|
||||
pool.ExpectQuery("name: GetClient :one").WithArgs(clientId).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "clientId", "canSync"}).
|
||||
AddRow(clientId, clientId, true),
|
||||
)
|
||||
pool.ExpectQuery("name: GetClient :one").WithArgs(clientId).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "name", "canSync"}).
|
||||
AddRow(clientId, "client_name", true),
|
||||
)
|
||||
|
||||
mockSQS.EXPECT().
|
||||
SendMessage(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
|
||||
return *in.QueueUrl == cfg.DocumentSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID.String())
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&sqs.SendMessageOutput{}, nil)
|
||||
|
||||
id, err := svc.Create(ctx, &Create{
|
||||
Hash: doc.Hash,
|
||||
Key: key,
|
||||
Bucket: bucket,
|
||||
Filename: nil, // No filename for this test
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, doc.ID, id)
|
||||
}
|
||||
|
||||
func TestGetCreateParams(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}),
|
||||
)
|
||||
|
||||
params, err := svc.getCreateParams(ctx, &Create{
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, &createDocumentParams{
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
}, params)
|
||||
}
|
||||
|
||||
func TestGetCreateParamsExisting(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(doc.ID),
|
||||
)
|
||||
|
||||
params, err := svc.getCreateParams(ctx, &Create{
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
Hash: doc.Hash,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
dbid := doc.ID
|
||||
assert.Equal(t, &createDocumentParams{
|
||||
ID: &dbid,
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
}, params)
|
||||
}
|
||||
|
||||
func TestGetCreateParamsCurrentDocErr(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).
|
||||
WillReturnError(errors.New("db err"))
|
||||
|
||||
_, err = svc.getCreateParams(ctx, &Create{
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
})
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSubmitCreate(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
// Set up mock S3 client for file size measurement
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
// Mock HeadObject to return file size
|
||||
fileSize := int64(2048)
|
||||
mockS3.EXPECT().
|
||||
HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == key.String()
|
||||
})).
|
||||
Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
|
||||
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(doc.ID),
|
||||
)
|
||||
pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(doc.ID, bucket, key.String()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
id, err := svc.submitCreate(ctx, &createDocumentParams{
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
Filename: nil, // No filename for this test
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, doc.ID, id)
|
||||
}
|
||||
|
||||
func TestSubmitCreateExists(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
// Set up mock S3 client for file size measurement
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
// Mock HeadObject to return file size (still called even for existing docs)
|
||||
fileSize := int64(4096)
|
||||
mockS3.EXPECT().
|
||||
HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == key.String()
|
||||
})).
|
||||
Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
|
||||
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(doc.ID, bucket, key.String()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
docid := doc.ID
|
||||
id, err := svc.submitCreate(ctx, &createDocumentParams{
|
||||
ID: &docid,
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, doc.ID, id)
|
||||
}
|
||||
|
||||
// TestSubmitCreate_HeadObjectError verifies that HeadObject errors are properly propagated.
|
||||
func TestSubmitCreate_HeadObjectError(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
// Set up mock S3 client that returns an error
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
// Mock HeadObject to return an error
|
||||
mockS3.EXPECT().
|
||||
HeadObject(mock.Anything, mock.Anything).
|
||||
Return(nil, errors.New("access denied"))
|
||||
|
||||
_, err = svc.submitCreate(ctx, &createDocumentParams{
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
})
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "failed to measure file size")
|
||||
assert.Contains(t, err.Error(), "access denied")
|
||||
}
|
||||
|
||||
// TestSubmitCreate_NoS3Client verifies that missing S3 client returns an error.
|
||||
func TestSubmitCreate_NoS3Client(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
// Intentionally not setting StoreClient
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: "hello",
|
||||
Hash: "example_hash",
|
||||
}
|
||||
bucket := "example_bucket"
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: doc.ClientID,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: uuid.New(),
|
||||
Location: objectstore.Import,
|
||||
}
|
||||
|
||||
_, err = svc.submitCreate(ctx, &createDocumentParams{
|
||||
Hash: doc.Hash,
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
})
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "S3 client not initialized")
|
||||
}
|
||||
|
||||
// TestMeasureFileSize_NilContentLength verifies handling of nil ContentLength.
|
||||
func TestMeasureFileSize_NilContentLength(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocInitConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
|
||||
// Set up mock S3 client that returns nil ContentLength
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
bucket := "example_bucket"
|
||||
key := "example_key"
|
||||
|
||||
mockS3.EXPECT().
|
||||
HeadObject(mock.Anything, mock.Anything).
|
||||
Return(&s3.HeadObjectOutput{ContentLength: nil}, nil)
|
||||
|
||||
_, err = svc.measureFileSize(ctx, bucket, key)
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "nil ContentLength")
|
||||
}
|
||||
|
||||
// TestMeasureFileSize_Integration tests the measureFileSize function using a real
|
||||
// S3 client (localstack) to verify that HeadObject correctly returns file sizes.
|
||||
// This is an integration test that requires Docker.
|
||||
|
||||
@@ -1,78 +0,0 @@
|
||||
package documentstore
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/queue/documentinit"
|
||||
queuemock "queryorchestration/mocks/queue"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/sqs"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
type ProcessConfig struct {
|
||||
serviceconfig.BaseConfig
|
||||
documentinit.DocInitConfig
|
||||
}
|
||||
|
||||
func TestProcess(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
cfg := &ProcessConfig{}
|
||||
mockSQS := queuemock.NewMockSQSClient(t)
|
||||
cfg.QueueClient = mockSQS
|
||||
cfg.DocInitURL = "docinit"
|
||||
|
||||
svc := New(cfg)
|
||||
|
||||
t.Run("invalid event", func(t *testing.T) {
|
||||
err := svc.Process(ctx, Params{
|
||||
Event: EventS3("invalid"),
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
t.Run("invalid key", func(t *testing.T) {
|
||||
err := svc.Process(ctx, Params{
|
||||
Event: EventS3ObjectCreatedPut,
|
||||
Key: "7db16095-9155-47d4-8004-b3b3ead93c83/invalid",
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
t.Run("document init", func(t *testing.T) {
|
||||
mockSQS.EXPECT().
|
||||
SendMessage(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
|
||||
return *in.QueueUrl == cfg.DocInitURL
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&sqs.SendMessageOutput{}, nil)
|
||||
|
||||
location := objectstore.BucketKey{
|
||||
ClientID: "7db16095-9155-47d4-8004-b3b3ead93c83",
|
||||
Location: objectstore.Import,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}
|
||||
err := svc.Process(ctx, Params{
|
||||
Event: EventS3ObjectCreatedPut,
|
||||
Key: location.String(),
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestIsSupportedEvent(t *testing.T) {
|
||||
svc := Service{}
|
||||
|
||||
t.Run("put", func(t *testing.T) {
|
||||
assert.True(t, svc.isSupportedEvent(EventS3ObjectCreatedPut))
|
||||
})
|
||||
t.Run("invalid", func(t *testing.T) {
|
||||
assert.False(t, svc.isSupportedEvent("invalid"))
|
||||
})
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"queryorchestration/internal/serviceconfig/queue/documentinit"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type DocInitConfig struct {
|
||||
@@ -18,3 +19,41 @@ func TestNewDocumentStoreService(t *testing.T) {
|
||||
svc := New(&DocInitConfig{})
|
||||
assert.NotNil(t, svc)
|
||||
}
|
||||
|
||||
func TestProcess_UnsupportedEvent(t *testing.T) {
|
||||
svc := New(&DocInitConfig{})
|
||||
|
||||
// Test unsupported event type - should return nil without error
|
||||
err := svc.Process(t.Context(), Params{
|
||||
Event: "ObjectRemoved:Delete",
|
||||
Key: "test-key",
|
||||
Bucket: "test-bucket",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestProcess_InvalidKey(t *testing.T) {
|
||||
svc := New(&DocInitConfig{})
|
||||
|
||||
// Test invalid key that can't be parsed - should return nil without error
|
||||
err := svc.Process(t.Context(), Params{
|
||||
Event: EventS3ObjectCreatedPut,
|
||||
Key: "invalid-key-without-proper-structure",
|
||||
Bucket: "test-bucket",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestProcess_NonImportLocation(t *testing.T) {
|
||||
svc := New(&DocInitConfig{})
|
||||
|
||||
// Test key with Export location - should return nil without error (unsupported key)
|
||||
// Key format: {client_id}/{location}/{date}/{filename}
|
||||
// Export location doesn't have a part number
|
||||
err := svc.Process(t.Context(), Params{
|
||||
Event: EventS3ObjectCreatedPut,
|
||||
Key: "testclient/export/20240101/2024-01-01T120000Z~testclient~export~550e8400-e29b-41d4-a716-446655440000.pdf",
|
||||
Bucket: "test-bucket",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
package documentsync_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"queryorchestration/internal/client"
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/document"
|
||||
documentsync "queryorchestration/internal/document/sync"
|
||||
queuemock "queryorchestration/mocks/queue"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/sqs"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pashagolub/pgxmock/v3"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
func TestSync(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
cfg := &DocSyncConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
mockSQS := queuemock.NewMockSQSClient(t)
|
||||
cfg.QueueClient = mockSQS
|
||||
cfg.DocumentCleanURL = "/i/am/here"
|
||||
|
||||
svc := documentsync.New(cfg, &documentsync.Services{
|
||||
Document: document.New(cfg),
|
||||
Client: client.New(cfg),
|
||||
})
|
||||
|
||||
j := client.Client{
|
||||
ID: "clientid",
|
||||
CanSync: true,
|
||||
}
|
||||
doc := document.DocumentSummary{
|
||||
ID: uuid.New(),
|
||||
ClientID: j.ID,
|
||||
Hash: "example_hash",
|
||||
}
|
||||
|
||||
pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(doc.ID).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "clientId", "hash"}).
|
||||
AddRow(doc.ID, doc.ClientID, doc.Hash),
|
||||
)
|
||||
pool.ExpectQuery("name: GetClient :one").WithArgs(j.ID).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "name", "canSync"}).
|
||||
AddRow(j.ID, "client_name", true),
|
||||
)
|
||||
|
||||
mockSQS.EXPECT().
|
||||
SendMessage(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
|
||||
return *in.QueueUrl == cfg.DocumentCleanURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID)
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&sqs.SendMessageOutput{}, nil)
|
||||
|
||||
err = svc.Sync(ctx, doc.ID)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
@@ -1,54 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
func (s *Service) hasCheckboxKeywords(text string) bool {
|
||||
conditions := []func(string) bool{
|
||||
s.hasExhibit,
|
||||
s.hasBenefit,
|
||||
s.hasPlans,
|
||||
s.hasServices,
|
||||
}
|
||||
|
||||
elementsContained := 0
|
||||
for _, hasCondition := range conditions {
|
||||
if hasCondition(text) {
|
||||
elementsContained++
|
||||
}
|
||||
}
|
||||
|
||||
return elementsContained >= 3
|
||||
}
|
||||
|
||||
func (s *Service) hasExhibit(text string) bool {
|
||||
return strings.Contains(text, "exhibit")
|
||||
}
|
||||
|
||||
func (s *Service) hasBenefit(text string) bool {
|
||||
keywords := []string{"benefit plan", "applicable benefit", "benefits"}
|
||||
|
||||
return orKeywords(text, keywords)
|
||||
}
|
||||
|
||||
func (s *Service) hasPlans(text string) bool {
|
||||
keywords := []string{"plan/program:", "plan(s)", "plans"}
|
||||
|
||||
return orKeywords(text, keywords)
|
||||
}
|
||||
|
||||
func (s *Service) hasServices(text string) bool {
|
||||
keywords := []string{"service:", "services", "service(s)"}
|
||||
return orKeywords(text, keywords)
|
||||
}
|
||||
|
||||
func orKeywords(text string, keywords []string) bool {
|
||||
for _, keyword := range keywords {
|
||||
if strings.Contains(text, keyword) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
@@ -1,57 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestHasCheckboxKeywords(t *testing.T) {
|
||||
svc := Service{}
|
||||
assert.False(t, svc.hasCheckboxKeywords(""))
|
||||
assert.False(t, svc.hasCheckboxKeywords("aaaaaa"))
|
||||
assert.False(t, svc.hasCheckboxKeywords("benefit plan"))
|
||||
assert.False(t, svc.hasCheckboxKeywords("benefit plans"))
|
||||
assert.True(t, svc.hasCheckboxKeywords("exhibit benefit plans"))
|
||||
}
|
||||
|
||||
func TestHasExhibit(t *testing.T) {
|
||||
svc := Service{}
|
||||
assert.False(t, svc.hasExhibit(""))
|
||||
assert.False(t, svc.hasExhibit("aaaaaa"))
|
||||
assert.True(t, svc.hasExhibit("exhibit"))
|
||||
}
|
||||
|
||||
func TestHasBenefit(t *testing.T) {
|
||||
svc := Service{}
|
||||
assert.False(t, svc.hasBenefit(""))
|
||||
assert.False(t, svc.hasBenefit("aaaaaa"))
|
||||
assert.False(t, svc.hasBenefit("benefit"))
|
||||
assert.True(t, svc.hasBenefit("benefit plan"))
|
||||
assert.True(t, svc.hasBenefit("benefit plan(s)"))
|
||||
assert.True(t, svc.hasBenefit("applicable benefit"))
|
||||
assert.True(t, svc.hasBenefit("benefits"))
|
||||
}
|
||||
|
||||
func TestHasPlans(t *testing.T) {
|
||||
svc := Service{}
|
||||
assert.False(t, svc.hasPlans(""))
|
||||
assert.False(t, svc.hasPlans("aaaaaa"))
|
||||
assert.True(t, svc.hasPlans("plan(s):"))
|
||||
assert.True(t, svc.hasPlans("plan(s)/program(s):"))
|
||||
assert.True(t, svc.hasPlans("plans:"))
|
||||
assert.True(t, svc.hasPlans("plan/program:"))
|
||||
assert.True(t, svc.hasPlans("plan(s)"))
|
||||
assert.True(t, svc.hasPlans("plans"))
|
||||
}
|
||||
|
||||
func TestHasServices(t *testing.T) {
|
||||
svc := Service{}
|
||||
assert.False(t, svc.hasServices(""))
|
||||
assert.False(t, svc.hasServices("aaaaaa"))
|
||||
assert.True(t, svc.hasServices("services:"))
|
||||
assert.True(t, svc.hasServices("service:"))
|
||||
assert.True(t, svc.hasServices("service(s):"))
|
||||
assert.True(t, svc.hasServices("services"))
|
||||
assert.True(t, svc.hasServices("service(s)"))
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (s *Service) GetDirectChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
|
||||
lines := []uuid.UUID{}
|
||||
found := map[uuid.UUID]*bool{}
|
||||
|
||||
for _, id := range s.getRelationshipIds(blockMap[id]) {
|
||||
if found[id] != nil && !*found[id] {
|
||||
continue
|
||||
}
|
||||
|
||||
lines = append(lines, id)
|
||||
f := true
|
||||
found[id] = &f
|
||||
|
||||
s.removeChildren(id, blockMap, &found)
|
||||
}
|
||||
|
||||
output := []uuid.UUID{}
|
||||
for _, id := range lines {
|
||||
if found[id] == nil || !*found[id] {
|
||||
continue
|
||||
}
|
||||
|
||||
output = append(output, id)
|
||||
}
|
||||
|
||||
return output
|
||||
}
|
||||
|
||||
func (s *Service) removeChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
|
||||
for _, childId := range s.getRelationshipIds(blockMap[id]) {
|
||||
f := false
|
||||
(*found)[childId] = &f
|
||||
s.removeChildren(childId, blockMap, found)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) GetBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
|
||||
blocks := make(map[uuid.UUID]types.Block, len(blocksList))
|
||||
for _, block := range blocksList {
|
||||
if block.Id == nil {
|
||||
continue
|
||||
}
|
||||
id, err := uuid.Parse(*block.Id)
|
||||
if err != nil {
|
||||
slog.Warn("error parsing block id for block map", "error", err, "id", *block.Id)
|
||||
continue
|
||||
}
|
||||
|
||||
blocks[id] = block
|
||||
}
|
||||
|
||||
return blocks
|
||||
}
|
||||
|
||||
func (s *Service) getRelationshipIds(block types.Block) []uuid.UUID {
|
||||
childIds := []uuid.UUID{}
|
||||
for _, relationship := range block.Relationships {
|
||||
for _, idStr := range relationship.Ids {
|
||||
id, err := uuid.Parse(idStr)
|
||||
if err != nil {
|
||||
slog.Error("invalid child uuid", "id", id, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
childIds = append(childIds, id)
|
||||
}
|
||||
}
|
||||
|
||||
return childIds
|
||||
}
|
||||
|
||||
func (s *Service) getRelationshipIdsByRel(block types.Block, relType types.RelationshipType) []uuid.UUID {
|
||||
childIds := []uuid.UUID{}
|
||||
for _, relationship := range block.Relationships {
|
||||
if relationship.Type != relType {
|
||||
continue
|
||||
}
|
||||
for _, idStr := range relationship.Ids {
|
||||
id, err := uuid.Parse(idStr)
|
||||
if err != nil {
|
||||
slog.Error("invalid child uuid", "id", id, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
childIds = append(childIds, id)
|
||||
}
|
||||
}
|
||||
|
||||
return childIds
|
||||
}
|
||||
@@ -1,62 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
documenttypes "queryorchestration/internal/document/types"
|
||||
)
|
||||
|
||||
func (s *Service) GetDocumentText(ctx context.Context, doc documenttypes.File) (io.Reader, error) {
|
||||
pagesNum, err := doc.GetPageCount(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
group := s.cfg.GetThreadPool().NewGroup()
|
||||
|
||||
textPages := make([]string, pagesNum)
|
||||
for i := range pagesNum {
|
||||
index := i
|
||||
group.SubmitErr(func() error {
|
||||
slog.Debug("extracting page text", "index", index)
|
||||
|
||||
text, err := s.getPageText(ctx, doc, index)
|
||||
if err != nil {
|
||||
slog.Error("error extracting page text", "index", index, "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
textPages[index] = text
|
||||
|
||||
slog.Info("extracted page text", "index", index)
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
err = group.Wait()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
text, err := s.mergeDocument(textPages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return text, nil
|
||||
}
|
||||
|
||||
func (s *Service) mergeDocument(pages []string) (io.Reader, error) {
|
||||
var builder strings.Builder
|
||||
for i, page := range pages {
|
||||
builder.WriteString(fmt.Sprintf("Start of Page No. = %d\n\n", i+1))
|
||||
builder.WriteString(page)
|
||||
builder.WriteString("\n\n")
|
||||
}
|
||||
return strings.NewReader(builder.String()), nil
|
||||
}
|
||||
@@ -1,406 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
documenttypes "queryorchestration/internal/document/types"
|
||||
textractmock "queryorchestration/mocks/textract"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract"
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetDocumentText(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
|
||||
t.Run("Hello World", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "helloWorld")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("Error Getting Textract Response", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
_, pdf, _, close := getFiles(t, "helloWorld")
|
||||
defer close()
|
||||
|
||||
mockTextract.EXPECT().
|
||||
AnalyzeDocument(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
||||
return true
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(nil, errors.New("no textract response"))
|
||||
|
||||
_, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.EqualError(t, err, "no textract response")
|
||||
})
|
||||
t.Run("merged_cell_table", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "merged_cell_table")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("table_check", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "table_check")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("table_select", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "table_select")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("multi_column", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "multi_column")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("table_check_row", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "table_check_row")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("table_aligned", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "table_aligned")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
|
||||
t.Run("cnc_01-06", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "cnc_01-06")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("hn_0109", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "hn_0109")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("hn_23-70", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "hn_23-70")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
t.Run("chc_1", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "chc_1")
|
||||
defer close()
|
||||
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
text, err := svc.GetDocumentText(ctx, pdf)
|
||||
require.NoError(t, err)
|
||||
|
||||
// With PDF processing, text output may differ from PNG-based expected files
|
||||
// Just verify we get some meaningful text output
|
||||
textBytes, err := io.ReadAll(text)
|
||||
require.NoError(t, err)
|
||||
textStr := string(textBytes)
|
||||
assert.NotEmpty(t, textStr, "Expected non-empty text output")
|
||||
assert.Contains(t, textStr, "Start of Page No.", "Expected page header in output")
|
||||
})
|
||||
}
|
||||
|
||||
func getFiles(t testing.TB, filename string) ([]map[string]*textract.AnalyzeDocumentOutput, *documenttypes.PDF, *os.File, func()) {
|
||||
root, err := os.OpenRoot("../../../assets")
|
||||
require.NoError(t, err)
|
||||
|
||||
textractFile, err := root.Open(fmt.Sprintf("textract/%s.gen", filename))
|
||||
require.NoError(t, err)
|
||||
defer textractFile.Close()
|
||||
|
||||
decoder := json.NewDecoder(textractFile)
|
||||
var out []map[string]*textract.AnalyzeDocumentOutput
|
||||
err = decoder.Decode(&out)
|
||||
require.NoError(t, err)
|
||||
|
||||
pdfFile, err := root.Open(fmt.Sprintf("original/%s.pdf", filename))
|
||||
require.NoError(t, err)
|
||||
|
||||
pdf, err := documenttypes.NewPDFFromReader(pdfFile)
|
||||
require.NoError(t, err)
|
||||
|
||||
txtFile, err := root.Open(fmt.Sprintf("text/%s.txt", filename))
|
||||
require.NoError(t, err)
|
||||
|
||||
return out, pdf, txtFile, func() {
|
||||
pdfFile.Close()
|
||||
txtFile.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractmock.MockTextractClient, pdf documenttypes.File, textractBaseOut []map[string]*textract.AnalyzeDocumentOutput) {
|
||||
// Set up flexible expectations that can handle variable number of calls per page
|
||||
// Use the actual textract data but allow for PDF vs PNG processing differences
|
||||
|
||||
// Find the first base response to use as template
|
||||
var baseResponse *textract.AnalyzeDocumentOutput
|
||||
var tableResponse *textract.AnalyzeDocumentOutput
|
||||
|
||||
for _, pageResult := range textractBaseOut {
|
||||
if baseResponse == nil && pageResult["base"] != nil {
|
||||
baseResponse = pageResult["base"]
|
||||
}
|
||||
if tableResponse == nil && pageResult["full_features"] != nil {
|
||||
tableResponse = pageResult["full_features"]
|
||||
}
|
||||
}
|
||||
|
||||
// Set up base layout call expectations - can be called multiple times
|
||||
if baseResponse != nil {
|
||||
mockTextract.EXPECT().
|
||||
AnalyzeDocument(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
||||
return len(in.Document.Bytes) > 0 && len(in.FeatureTypes) > 0 && in.FeatureTypes[0] == types.FeatureTypeLayout
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(baseResponse, nil).Maybe()
|
||||
}
|
||||
|
||||
// Set up table/forms call expectations if they exist - can be called multiple times
|
||||
if tableResponse != nil {
|
||||
mockTextract.EXPECT().
|
||||
AnalyzeDocument(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
||||
if len(in.Document.Bytes) == 0 || len(in.FeatureTypes) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, feature := range in.FeatureTypes {
|
||||
if feature == types.FeatureTypeForms || feature == types.FeatureTypeTables {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(tableResponse, nil).Maybe()
|
||||
}
|
||||
}
|
||||
@@ -1,48 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
|
||||
documenttypes "queryorchestration/internal/document/types"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (s *Service) executeExtraction(ctx context.Context, documentId uuid.UUID) error {
|
||||
clean, err := s.cfg.GetDBQueries().GetCleanEntryByDocId(ctx, documentId)
|
||||
if err != nil {
|
||||
return err
|
||||
} else if clean.Fail.Valid {
|
||||
slog.Error("document has no clean document", "id", documentId)
|
||||
return nil
|
||||
}
|
||||
|
||||
key, err := objectstore.ParseBucketKey(*clean.Key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
doc, err := documenttypes.GetFile(ctx, s.cfg, documenttypes.GetFileParams{
|
||||
Bucket: *clean.Bucket,
|
||||
Key: key,
|
||||
Hash: *clean.Hash,
|
||||
Mimetype: documenttypes.ParseDBNullMimeType(clean.Mimetype),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
text, err := s.GetDocumentText(ctx, doc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
err = s.storeProcess(ctx, text, clean)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,89 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
textractmock "queryorchestration/mocks/textract"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pashagolub/pgxmock/v3"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestExecuteExtract(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
|
||||
cfg := &DocTextConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
mockStore := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockStore
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
cleanId := uuid.New()
|
||||
docId := uuid.New()
|
||||
bucket := "bucket"
|
||||
// Hash will be different with PDF vs PNG processing - use actual computed hash
|
||||
// We'll capture the computed hash and use it in expectations
|
||||
hash := `"computed-during-test"`
|
||||
textId := uuid.New()
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: "aa",
|
||||
Location: objectstore.Import,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: cleanId,
|
||||
}
|
||||
keyStr := key.String()
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "helloWorld")
|
||||
defer close()
|
||||
|
||||
mimetype := repository.NullCleanmimetype{
|
||||
Valid: true,
|
||||
Cleanmimetype: repository.CleanmimetypeApplicationPdf,
|
||||
}
|
||||
pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(docId).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}).
|
||||
AddRow(cleanId, docId, &bucket, &keyStr, int64(123), &hash, mimetype, nil, key.ClientID),
|
||||
)
|
||||
bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
|
||||
mockStore.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == keyStr
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: bodyMsg,
|
||||
}, nil)
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, pgxmock.AnyArg()).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(textId),
|
||||
)
|
||||
pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
err = svc.executeExtraction(ctx, docId)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
@@ -1,30 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log/slog"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Extract extracts text from a document using AWS Textract.
|
||||
// This is the terminal step in the document processing pipeline.
|
||||
// Query processing has been removed - pipeline ends after text extraction.
|
||||
func (s *Service) Extract(ctx context.Context, documentId uuid.UUID) error {
|
||||
slog.Debug("extracting document text", "id", documentId.String())
|
||||
|
||||
isextracted, err := s.cfg.GetDBQueries().IsDocumentTextExtracted(ctx, documentId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if !isextracted {
|
||||
err = s.executeExtraction(ctx, documentId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Pipeline terminates here - query processing has been removed
|
||||
return nil
|
||||
}
|
||||
@@ -1,114 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/aws"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/textract"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
textractmock "queryorchestration/mocks/textract"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pashagolub/pgxmock/v3"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type DocTextConfig struct {
|
||||
aws.AWSConfig
|
||||
serviceconfig.BaseConfig
|
||||
textract.TextractConfig
|
||||
objectstore.ObjectStoreConfig
|
||||
}
|
||||
|
||||
func TestCreate(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
|
||||
cfg := &DocTextConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
mockStore := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockStore
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
docId := uuid.New()
|
||||
t.Run("is not extracted", func(t *testing.T) {
|
||||
cleanId := uuid.New()
|
||||
bucket := "bucket"
|
||||
// Hash will be different with PDF vs PNG processing - use actual computed hash
|
||||
// We'll capture the computed hash and use it in expectations
|
||||
hash := `"computed-during-test"`
|
||||
textId := uuid.New()
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: "aa",
|
||||
Location: objectstore.Import,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: cleanId,
|
||||
}
|
||||
keyStr := key.String()
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "helloWorld")
|
||||
defer close()
|
||||
|
||||
mimetype := repository.NullCleanmimetype{
|
||||
Valid: true,
|
||||
Cleanmimetype: repository.CleanmimetypeApplicationPdf,
|
||||
}
|
||||
pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(docId).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"isextracted"}).
|
||||
AddRow(false),
|
||||
)
|
||||
pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(docId).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}).
|
||||
AddRow(cleanId, docId, &bucket, &keyStr, int64(123), &hash, mimetype, nil, key.ClientID),
|
||||
)
|
||||
bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
|
||||
mockStore.EXPECT().
|
||||
GetObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
|
||||
return *in.Bucket == bucket && *in.Key == keyStr
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.GetObjectOutput{
|
||||
Body: bodyMsg,
|
||||
}, nil)
|
||||
textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
|
||||
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, pgxmock.AnyArg()).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(textId),
|
||||
)
|
||||
pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
err = svc.Extract(ctx, docId)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
t.Run("is extracted", func(t *testing.T) {
|
||||
pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(docId).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"isextracted"}).
|
||||
AddRow(true),
|
||||
)
|
||||
|
||||
err = svc.Extract(ctx, docId)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
@@ -1,113 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (s *Service) getKeyValueSets(id uuid.UUID, elements PageElements) (map[uuid.UUID]uuid.UUID, error) {
|
||||
block := elements.BlockMap[id]
|
||||
keyValuePairs := map[uuid.UUID]bool{}
|
||||
lines := map[uuid.UUID]bool{}
|
||||
lineToKey := map[uuid.UUID]uuid.UUID{}
|
||||
|
||||
for _, lineId := range s.getRelationshipIds(block) {
|
||||
for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
|
||||
for _, b := range elements.BlockMap {
|
||||
for _, childId := range s.getRelationshipIds(b) {
|
||||
if wordId != childId || b.BlockType != types.BlockTypeKeyValueSet {
|
||||
continue
|
||||
}
|
||||
|
||||
id, err := uuid.Parse(*b.Id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
lines[lineId] = true
|
||||
keyValuePairs[id] = true
|
||||
lineToKey[lineId] = id
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, lineId := range s.getRelationshipIds(block) {
|
||||
if lines[lineId] {
|
||||
continue
|
||||
}
|
||||
for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
|
||||
for _, b := range elements.BlockMap {
|
||||
for _, childId := range s.getRelationshipIds(b) {
|
||||
id, err := uuid.Parse(*b.Id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if wordId != childId || keyValuePairs[id] {
|
||||
continue
|
||||
}
|
||||
keyValuePairs[id] = true
|
||||
lineToKey[lineId] = id
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return lineToKey, nil
|
||||
}
|
||||
|
||||
func (s *Service) addKeyValue(id uuid.UUID, elements PageElements, builder *strings.Builder) error {
|
||||
lineToKey, err := s.getKeyValueSets(id, elements)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, lineId := range s.getRelationshipIds(elements.BlockMap[id]) {
|
||||
keysBlock := lineToKey[lineId]
|
||||
needsSelection := false
|
||||
isSelected := false
|
||||
|
||||
for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeValue) {
|
||||
for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[id], types.RelationshipTypeChild) {
|
||||
keyBlock := elements.BlockMap[id]
|
||||
if keyBlock.BlockType != types.BlockTypeSelectionElement {
|
||||
continue
|
||||
}
|
||||
needsSelection = true
|
||||
if keyBlock.SelectionStatus == types.SelectionStatusSelected {
|
||||
isSelected = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !needsSelection {
|
||||
builder.WriteString(fmt.Sprintf("%s\n", *elements.BlockMap[lineId].Text))
|
||||
continue
|
||||
} else if !isSelected {
|
||||
continue
|
||||
}
|
||||
|
||||
var key strings.Builder
|
||||
for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeChild) {
|
||||
keyBlock := elements.BlockMap[id]
|
||||
if keyBlock.BlockType != types.BlockTypeWord {
|
||||
slog.Info("uncaught element in key value", "type", keyBlock.BlockType)
|
||||
continue
|
||||
}
|
||||
if key.String() != "" {
|
||||
key.WriteRune(' ')
|
||||
}
|
||||
key.WriteString(*keyBlock.Text)
|
||||
}
|
||||
|
||||
builder.WriteString(fmt.Sprintf("%s\n", key.String()))
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,226 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
documenttypes "queryorchestration/internal/document/types"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract"
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (s *Service) getPageText(ctx context.Context, document documenttypes.File, index int) (string, error) {
|
||||
res, err := s.GetBasePage(ctx, document, index)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
features, err := s.ListRequiredFeatures(ctx, res.Blocks, index)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
blocks := res.Blocks
|
||||
if len(features) > 0 {
|
||||
res, err := s.GetPageWithFeatures(ctx, document, index, features)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
blocks = res.Blocks
|
||||
}
|
||||
|
||||
elements, err := s.getPageElements(ctx, document, index, blocks)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return s.buildPageText(ctx, elements, index)
|
||||
}
|
||||
|
||||
type PageElements struct {
|
||||
Page uuid.UUID
|
||||
BlockMap map[uuid.UUID]types.Block
|
||||
DirectChildren []uuid.UUID
|
||||
Tables []uuid.UUID
|
||||
}
|
||||
|
||||
func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, blocks []types.Block) (PageElements, error) {
|
||||
elements := PageElements{}
|
||||
|
||||
pageBlock, err := s.GetPageBlock(blocks)
|
||||
if err != nil {
|
||||
return PageElements{}, err
|
||||
}
|
||||
|
||||
elements.Page = pageBlock
|
||||
elements.BlockMap = s.GetBlockMap(blocks)
|
||||
elements.DirectChildren = s.GetDirectChildren(pageBlock, elements.BlockMap)
|
||||
|
||||
for _, block := range blocks {
|
||||
if block.BlockType != types.BlockTypeTable {
|
||||
continue
|
||||
}
|
||||
id, err := uuid.Parse(*block.Id)
|
||||
if err != nil {
|
||||
return PageElements{}, err
|
||||
}
|
||||
|
||||
elements.Tables = append(elements.Tables, id)
|
||||
}
|
||||
|
||||
return elements, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListRequiredFeatures(ctx context.Context, blocks []types.Block, index int) ([]types.FeatureType, error) {
|
||||
features := []types.FeatureType{}
|
||||
for _, block := range blocks {
|
||||
switch block.BlockType {
|
||||
case types.BlockTypeLayoutTable:
|
||||
features = append(features, types.FeatureTypeTables)
|
||||
case types.BlockTypeLayoutKeyValue:
|
||||
features = append(features, types.FeatureTypeForms)
|
||||
}
|
||||
}
|
||||
|
||||
if len(features) > 0 {
|
||||
features = append(features, types.FeatureTypeLayout)
|
||||
features = append(features, types.FeatureTypeSignatures)
|
||||
}
|
||||
|
||||
return features, nil
|
||||
}
|
||||
|
||||
type buildParams struct {
|
||||
tableIndex int
|
||||
}
|
||||
|
||||
func (s *Service) buildPageText(ctx context.Context, elements PageElements, index int) (string, error) {
|
||||
var builder strings.Builder
|
||||
params := buildParams{}
|
||||
|
||||
for _, id := range elements.DirectChildren {
|
||||
err := s.addComponentText(ctx, id, elements, ¶ms, &builder)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
for params.tableIndex < len(elements.Tables) {
|
||||
err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, &builder)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
params.tableIndex++
|
||||
}
|
||||
|
||||
numOfSignatures := 0
|
||||
for _, block := range elements.BlockMap {
|
||||
if block.BlockType == types.BlockTypeSignature {
|
||||
numOfSignatures++
|
||||
}
|
||||
}
|
||||
|
||||
if numOfSignatures > 0 {
|
||||
builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", numOfSignatures))
|
||||
}
|
||||
|
||||
return builder.String(), nil
|
||||
}
|
||||
|
||||
func (s *Service) addComponentText(ctx context.Context, id uuid.UUID, elements PageElements, params *buildParams, builder *strings.Builder) error {
|
||||
block := elements.BlockMap[id]
|
||||
switch block.BlockType {
|
||||
case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure:
|
||||
for _, id := range s.getRelationshipIds(block) {
|
||||
err := s.addComponentText(ctx, id, elements, params, builder)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case types.BlockTypeLine:
|
||||
_, err := builder.WriteString(fmt.Sprintf("%s\n", *block.Text))
|
||||
if err != nil {
|
||||
slog.Error("unable to write string", "error", err)
|
||||
return err
|
||||
}
|
||||
case types.BlockTypeWord:
|
||||
_, err := builder.WriteString(fmt.Sprintf(" %s", *block.Text))
|
||||
if err != nil {
|
||||
slog.Error("unable to write string", "error", err)
|
||||
return err
|
||||
}
|
||||
case types.BlockTypeLayoutTable:
|
||||
if params.tableIndex < len(elements.Tables) {
|
||||
err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, builder)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
for _, child := range s.getRelationshipIds(block) {
|
||||
err := s.addComponentText(ctx, child, elements, params, builder)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
params.tableIndex++
|
||||
case types.BlockTypeLayoutKeyValue:
|
||||
err := s.addKeyValue(id, elements, builder)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
default:
|
||||
slog.Debug("uncaught component", "type", block.BlockType)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index int) (textract.AnalyzeDocumentOutput, error) {
|
||||
return s.GetPageWithFeatures(ctx, doc, index, []types.FeatureType{
|
||||
types.FeatureTypeLayout,
|
||||
types.FeatureTypeSignatures,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) GetPageWithFeatures(ctx context.Context, doc documenttypes.File, index int, features []types.FeatureType) (textract.AnalyzeDocumentOutput, error) {
|
||||
pageBytes, err := doc.GetPage(ctx, index)
|
||||
if err != nil {
|
||||
return textract.AnalyzeDocumentOutput{}, err
|
||||
}
|
||||
|
||||
res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
|
||||
Document: &types.Document{
|
||||
Bytes: pageBytes,
|
||||
},
|
||||
FeatureTypes: features,
|
||||
})
|
||||
if err != nil {
|
||||
return textract.AnalyzeDocumentOutput{}, err
|
||||
}
|
||||
|
||||
return *res, err
|
||||
}
|
||||
|
||||
func (s *Service) GetPageBlock(blocks []types.Block) (uuid.UUID, error) {
|
||||
var pageBlock *types.Block
|
||||
for _, block := range blocks {
|
||||
if block.BlockType == types.BlockTypePage {
|
||||
pageBlock = &block
|
||||
break
|
||||
}
|
||||
}
|
||||
if pageBlock == nil {
|
||||
return uuid.Nil, errors.New("no page block found")
|
||||
}
|
||||
|
||||
return uuid.Parse(*pageBlock.Id)
|
||||
}
|
||||
@@ -1,329 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
textractmock "queryorchestration/mocks/textract"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract"
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAddComponentText(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
svc := Service{}
|
||||
|
||||
t.Run("table with index out of bounds", func(t *testing.T) {
|
||||
// Test the else branch when tableIndex >= len(elements.Tables)
|
||||
blockId := uuid.New()
|
||||
childId1 := uuid.New()
|
||||
childId2 := uuid.New()
|
||||
|
||||
elements := PageElements{
|
||||
BlockMap: map[uuid.UUID]types.Block{
|
||||
blockId: {
|
||||
BlockType: types.BlockTypeLayoutTable,
|
||||
Relationships: []types.Relationship{
|
||||
{
|
||||
Type: types.RelationshipTypeChild,
|
||||
Ids: []string{childId1.String(), childId2.String()},
|
||||
},
|
||||
},
|
||||
},
|
||||
childId1: {
|
||||
BlockType: types.BlockTypeLine,
|
||||
Text: &[]string{"Table child text 1"}[0],
|
||||
},
|
||||
childId2: {
|
||||
BlockType: types.BlockTypeLine,
|
||||
Text: &[]string{"Table child text 2"}[0],
|
||||
},
|
||||
},
|
||||
Tables: []uuid.UUID{}, // Empty tables array to trigger else branch
|
||||
}
|
||||
|
||||
params := &buildParams{
|
||||
tableIndex: 0, // This will be >= len(elements.Tables) which is 0
|
||||
}
|
||||
|
||||
var builder strings.Builder
|
||||
err := svc.addComponentText(ctx, blockId, elements, params, &builder)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, builder.String(), "Table child text 1")
|
||||
assert.Contains(t, builder.String(), "Table child text 2")
|
||||
assert.Equal(t, 1, params.tableIndex) // Should be incremented
|
||||
})
|
||||
|
||||
t.Run("default block type", func(t *testing.T) {
|
||||
// Test the default case for unhandled block types
|
||||
blockId := uuid.New()
|
||||
|
||||
elements := PageElements{
|
||||
BlockMap: map[uuid.UUID]types.Block{
|
||||
blockId: {
|
||||
BlockType: types.BlockTypeQuery, // An unhandled type to trigger default case
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
params := &buildParams{}
|
||||
var builder strings.Builder
|
||||
|
||||
// The default case just logs debug and returns nil
|
||||
err := svc.addComponentText(ctx, blockId, elements, params, &builder)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, builder.String())
|
||||
})
|
||||
|
||||
t.Run("word block type", func(t *testing.T) {
|
||||
// Test BlockTypeWord path
|
||||
blockId := uuid.New()
|
||||
|
||||
elements := PageElements{
|
||||
BlockMap: map[uuid.UUID]types.Block{
|
||||
blockId: {
|
||||
BlockType: types.BlockTypeWord,
|
||||
Text: &[]string{"TestWord"}[0],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
params := &buildParams{}
|
||||
var builder strings.Builder
|
||||
|
||||
err := svc.addComponentText(ctx, blockId, elements, params, &builder)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, " TestWord", builder.String())
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetBasePage(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
cfg := DocTextConfig{}
|
||||
svc := Service{
|
||||
cfg: &cfg,
|
||||
}
|
||||
|
||||
mockTextract := textractmock.NewMockTextractClient(t)
|
||||
cfg.TextractClient = mockTextract
|
||||
|
||||
textractBaseOut, pdf, _, close := getFiles(t, "helloWorld")
|
||||
defer close()
|
||||
|
||||
base := textractBaseOut[0]["base"]
|
||||
|
||||
mockTextract.EXPECT().
|
||||
AnalyzeDocument(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
||||
return len(in.Document.Bytes) > 0 && in.FeatureTypes[0] == types.FeatureTypeLayout
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(base, nil).Once()
|
||||
|
||||
text, err := svc.GetBasePage(ctx, pdf, 0)
|
||||
require.NoError(t, err)
|
||||
assert.EqualExportedValues(t, *base, text)
|
||||
}
|
||||
|
||||
func TestGetBlockMap(t *testing.T) {
|
||||
svc := Service{}
|
||||
t.Run("0 blocks", func(t *testing.T) {
|
||||
blocks := []types.Block{}
|
||||
outBlocks := svc.GetBlockMap(blocks)
|
||||
assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
|
||||
})
|
||||
t.Run("nil id blocks", func(t *testing.T) {
|
||||
blocks := []types.Block{
|
||||
{},
|
||||
}
|
||||
outBlocks := svc.GetBlockMap(blocks)
|
||||
assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
|
||||
})
|
||||
t.Run("1 block", func(t *testing.T) {
|
||||
blockId := uuid.New()
|
||||
blockIdStr := blockId.String()
|
||||
blocks := []types.Block{
|
||||
{
|
||||
Id: &blockIdStr,
|
||||
},
|
||||
}
|
||||
outBlocks := svc.GetBlockMap(blocks)
|
||||
assert.Equal(t, map[uuid.UUID]types.Block{
|
||||
blockId: {
|
||||
Id: &blockIdStr,
|
||||
},
|
||||
}, outBlocks)
|
||||
})
|
||||
t.Run("invalid block id", func(t *testing.T) {
|
||||
blockIdStr := "invalid_block_id"
|
||||
blocks := []types.Block{
|
||||
{
|
||||
Id: &blockIdStr,
|
||||
},
|
||||
}
|
||||
outBlocks := svc.GetBlockMap(blocks)
|
||||
assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
|
||||
})
|
||||
t.Run("2 blocks", func(t *testing.T) {
|
||||
blockOneId := uuid.New()
|
||||
blockOneIdStr := blockOneId.String()
|
||||
blockTwoId := uuid.New()
|
||||
blockTwoIdStr := blockTwoId.String()
|
||||
blocks := []types.Block{
|
||||
{
|
||||
Id: &blockOneIdStr,
|
||||
},
|
||||
{
|
||||
Id: &blockTwoIdStr,
|
||||
},
|
||||
}
|
||||
outBlocks := svc.GetBlockMap(blocks)
|
||||
assert.Equal(t, map[uuid.UUID]types.Block{
|
||||
blockOneId: {
|
||||
Id: &blockOneIdStr,
|
||||
},
|
||||
blockTwoId: {
|
||||
Id: &blockTwoIdStr,
|
||||
},
|
||||
}, outBlocks)
|
||||
})
|
||||
}
|
||||
|
||||
func TestMergeDocument(t *testing.T) {
|
||||
svc := Service{}
|
||||
t.Run("0 pages", func(t *testing.T) {
|
||||
pages := []string{}
|
||||
out, err := svc.mergeDocument(pages)
|
||||
require.NoError(t, err)
|
||||
outStr, err := io.ReadAll(out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ``, string(outStr))
|
||||
})
|
||||
t.Run("1 page", func(t *testing.T) {
|
||||
pages := []string{
|
||||
"page one",
|
||||
}
|
||||
out, err := svc.mergeDocument(pages)
|
||||
require.NoError(t, err)
|
||||
outStr, err := io.ReadAll(out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `Start of Page No. = 1
|
||||
|
||||
page one
|
||||
|
||||
`, string(outStr))
|
||||
})
|
||||
t.Run("2 pages", func(t *testing.T) {
|
||||
pages := []string{
|
||||
"page one",
|
||||
"page two",
|
||||
}
|
||||
out, err := svc.mergeDocument(pages)
|
||||
require.NoError(t, err)
|
||||
outStr, err := io.ReadAll(out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `Start of Page No. = 1
|
||||
|
||||
page one
|
||||
|
||||
Start of Page No. = 2
|
||||
|
||||
page two
|
||||
|
||||
`, string(outStr))
|
||||
})
|
||||
t.Run("3 pages", func(t *testing.T) {
|
||||
pages := []string{
|
||||
"page one",
|
||||
"page two",
|
||||
"page three",
|
||||
}
|
||||
out, err := svc.mergeDocument(pages)
|
||||
require.NoError(t, err)
|
||||
outStr, err := io.ReadAll(out)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `Start of Page No. = 1
|
||||
|
||||
page one
|
||||
|
||||
Start of Page No. = 2
|
||||
|
||||
page two
|
||||
|
||||
Start of Page No. = 3
|
||||
|
||||
page three
|
||||
|
||||
`, string(outStr))
|
||||
})
|
||||
}
|
||||
|
||||
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`
|
||||
@@ -1,23 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/textract"
|
||||
)
|
||||
|
||||
type ConfigProvider interface {
|
||||
serviceconfig.ConfigProvider
|
||||
textract.ConfigProvider
|
||||
objectstore.ConfigProvider
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
cfg ConfigProvider
|
||||
}
|
||||
|
||||
func New(cfg ConfigProvider) *Service {
|
||||
return &Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
package documenttext_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
documenttext "queryorchestration/internal/document/text"
|
||||
"queryorchestration/internal/serviceconfig"
|
||||
"queryorchestration/internal/serviceconfig/aws"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
"queryorchestration/internal/serviceconfig/textract"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
type DocTextConfig struct {
|
||||
aws.AWSConfig
|
||||
serviceconfig.BaseConfig
|
||||
textract.TextractConfig
|
||||
objectstore.ObjectStoreConfig
|
||||
}
|
||||
|
||||
func TestService(t *testing.T) {
|
||||
cfg := &DocTextConfig{}
|
||||
svc := documenttext.New(cfg)
|
||||
assert.NotNil(t, svc)
|
||||
}
|
||||
@@ -1,118 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/serviceconfig/build"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
)
|
||||
|
||||
func (s *Service) storeProcess(ctx context.Context, text io.Reader, clean *repository.Currentcleanentry) error {
|
||||
version := build.GetVersionUnixTimestamp()
|
||||
|
||||
hashReader, storeReader, err := duplicateReader(text)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
hash, err := s.cfg.CalculateETag(ctx, hashReader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
|
||||
existingId, err := q.GetDocumentTextExtractionByHash(ctx, &repository.GetDocumentTextExtractionByHashParams{
|
||||
Cleanentryid: clean.ID,
|
||||
Hash: hash,
|
||||
})
|
||||
if err != nil && !errors.Is(err, sql.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
|
||||
var textId uuid.UUID
|
||||
if existingId == uuid.Nil || errors.Is(err, sql.ErrNoRows) {
|
||||
createdAt := time.Now().UTC()
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: clean.Clientid,
|
||||
Location: objectstore.Text,
|
||||
CreatedAt: createdAt,
|
||||
EntityID: clean.ID,
|
||||
}
|
||||
|
||||
currentPart, err := q.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{
|
||||
Clientid: &key.ClientID,
|
||||
Querydate: pgtype.Timestamptz{
|
||||
Time: key.CreatedAt,
|
||||
Valid: true,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
part := s.cfg.GetDirectoryPart(currentPart.Part, currentPart.Count)
|
||||
key.Part = &part
|
||||
|
||||
keyStr := key.String()
|
||||
|
||||
textId, err = q.AddDocumentText(ctx, &repository.AddDocumentTextParams{
|
||||
Cleanid: clean.ID,
|
||||
Bucket: *clean.Bucket,
|
||||
Key: keyStr,
|
||||
Hash: hash,
|
||||
Part: *key.Part,
|
||||
Createdat: pgtype.Timestamp{
|
||||
Time: key.CreatedAt,
|
||||
Valid: true,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = s.cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: clean.Bucket,
|
||||
Key: &keyStr,
|
||||
Body: storeReader,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
} else {
|
||||
textId = existingId
|
||||
}
|
||||
|
||||
err = q.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{
|
||||
Textid: textId,
|
||||
Version: version,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func duplicateReader(reader io.Reader) (io.Reader, io.Reader, error) {
|
||||
content, err := io.ReadAll(reader)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
reader1 := bytes.NewReader(content)
|
||||
reader2 := bytes.NewReader(content)
|
||||
|
||||
return reader1, reader2, nil
|
||||
}
|
||||
@@ -1,121 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/serviceconfig/objectstore"
|
||||
objectstoremock "queryorchestration/mocks/objectstore"
|
||||
|
||||
"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"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestStoreExtract(t *testing.T) {
|
||||
ctx := t.Context()
|
||||
pool, err := pgxmock.NewPool()
|
||||
require.NoError(t, err)
|
||||
|
||||
cfg := &DocTextConfig{}
|
||||
cfg.DBPool = pool
|
||||
cfg.DBQueries = repository.New(pool)
|
||||
mockS3 := objectstoremock.NewMockS3Client(t)
|
||||
cfg.StoreClient = mockS3
|
||||
|
||||
svc := Service{
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
cleanId := uuid.New()
|
||||
textId := uuid.New()
|
||||
clientId := "hello"
|
||||
bucket := "bucket"
|
||||
|
||||
t.Run("exists", func(t *testing.T) {
|
||||
hash := `"5d41402abc4b2a76b9719d911017c592"`
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(textId),
|
||||
)
|
||||
pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
text := strings.NewReader("hello")
|
||||
err = svc.storeProcess(ctx, text, &repository.Currentcleanentry{
|
||||
ID: cleanId,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
t.Run("not exists", func(t *testing.T) {
|
||||
part := uint16(0)
|
||||
key := objectstore.BucketKey{
|
||||
ClientID: clientId,
|
||||
Location: objectstore.Text,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
EntityID: cleanId,
|
||||
Part: &part,
|
||||
}
|
||||
keyStr := key.String()
|
||||
hash := `"09644372e99020106946045c6fd2d70b"`
|
||||
pool.ExpectBegin()
|
||||
pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
pool.ExpectQuery("name: GetTextOutCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"part", "count"}).
|
||||
AddRow(0, 2),
|
||||
)
|
||||
pool.ExpectQuery("name: AddDocumentText :one").WithArgs(cleanId, bucket, key.String(), hash, pgxmock.AnyArg(), uint16(0)).
|
||||
WillReturnRows(
|
||||
pgxmock.NewRows([]string{"id"}).
|
||||
AddRow(textId),
|
||||
)
|
||||
|
||||
mockS3.EXPECT().
|
||||
PutObject(
|
||||
mock.Anything,
|
||||
mock.MatchedBy(func(in *s3.PutObjectInput) bool {
|
||||
bod, err := io.ReadAll(in.Body)
|
||||
require.NoError(t, err)
|
||||
return string(bod) == "byebye" && *in.Key == key.String() && *in.Bucket == bucket
|
||||
}),
|
||||
mock.Anything,
|
||||
).
|
||||
Return(&s3.PutObjectOutput{}, nil)
|
||||
|
||||
pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
|
||||
WillReturnResult(pgxmock.NewResult("", 1))
|
||||
pool.ExpectCommit()
|
||||
|
||||
text := strings.NewReader("byebye")
|
||||
err = svc.storeProcess(ctx, text, &repository.Currentcleanentry{
|
||||
ID: cleanId,
|
||||
Clientid: clientId,
|
||||
Bucket: &bucket,
|
||||
Key: &keyStr,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestDuplicateReader(t *testing.T) {
|
||||
original := strings.NewReader("hello")
|
||||
readerOne, readerTwo, err := duplicateReader(original)
|
||||
require.NoError(t, err)
|
||||
contentOne, err := io.ReadAll(readerOne)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []byte("hello"), contentOne)
|
||||
contentTwo, err := io.ReadAll(readerTwo)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []byte("hello"), contentTwo)
|
||||
}
|
||||
@@ -1,163 +0,0 @@
|
||||
package documenttext
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"strings"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/textract/types"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
type cell struct {
|
||||
value string
|
||||
row int
|
||||
column int
|
||||
isSelected bool
|
||||
}
|
||||
|
||||
func (s *Service) addTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
|
||||
params := tableParams{
|
||||
maxRow: 0,
|
||||
maxColumn: 0,
|
||||
table: []cell{},
|
||||
}
|
||||
|
||||
for _, uid := range s.GetDirectChildren(id, blockMap) {
|
||||
s.addComponentTable(ctx, uid, blockMap, ¶ms)
|
||||
}
|
||||
|
||||
invalidRows := make([]bool, params.maxRow+1)
|
||||
for _, cell := range params.table {
|
||||
if !cell.isSelected {
|
||||
invalidRows[cell.row] = true
|
||||
}
|
||||
}
|
||||
|
||||
tableArr := make([][]string, params.maxRow+1)
|
||||
for i := range tableArr {
|
||||
tableArr[i] = make([]string, params.maxColumn+1)
|
||||
}
|
||||
|
||||
for _, cell := range params.table {
|
||||
tableArr[cell.row][cell.column] = cell.value
|
||||
}
|
||||
|
||||
finalTable := [][]string{}
|
||||
for i, row := range tableArr {
|
||||
if !invalidRows[i] {
|
||||
finalTable = append(finalTable, row)
|
||||
}
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(finalTable)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
builder.WriteString("-------Table Start--------\n")
|
||||
builder.WriteString(string(jsonData))
|
||||
builder.WriteString("\n-------Table End--------\n")
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
type tableParams struct {
|
||||
maxRow int
|
||||
maxColumn int
|
||||
table []cell
|
||||
}
|
||||
|
||||
func (s *Service) addComponentTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, params *tableParams) {
|
||||
block := blockMap[id]
|
||||
|
||||
switch block.BlockType {
|
||||
case types.BlockTypeMergedCell:
|
||||
content := ""
|
||||
addRow := true
|
||||
for _, uid := range s.getRelationshipIds(block) {
|
||||
if blockMap[uid].BlockType != types.BlockTypeCell {
|
||||
slog.Warn("not a cell", "block_type", blockMap[uid].BlockType)
|
||||
continue
|
||||
}
|
||||
content, addRow = s.getCellContent(uid, blockMap)
|
||||
if content != "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
startColumnIndex := int(*block.ColumnIndex - 1)
|
||||
startRowIndex := int(*block.RowIndex - 1)
|
||||
endColumnIndex := startColumnIndex + int(*block.ColumnSpan) - 1
|
||||
endRowIndex := startRowIndex + int(*block.RowSpan) - 1
|
||||
for row := startRowIndex; row <= endRowIndex; row++ {
|
||||
for column := startColumnIndex; column <= endColumnIndex; column++ {
|
||||
s.addCellToTable(
|
||||
row,
|
||||
column,
|
||||
content,
|
||||
addRow,
|
||||
params,
|
||||
)
|
||||
}
|
||||
}
|
||||
case types.BlockTypeCell:
|
||||
content, addRow := s.getCellContent(id, blockMap)
|
||||
s.addCellToTable(
|
||||
int(*block.RowIndex)-1,
|
||||
int(*block.ColumnIndex)-1,
|
||||
content,
|
||||
addRow,
|
||||
params,
|
||||
)
|
||||
default:
|
||||
slog.Info("unhandled block type in table", "type", block.BlockType)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) getCellContent(id uuid.UUID, blockMap map[uuid.UUID]types.Block) (string, bool) {
|
||||
var str strings.Builder
|
||||
var isSelected *bool
|
||||
for _, uid := range s.getRelationshipIds(blockMap[id]) {
|
||||
if blockMap[uid].BlockType == types.BlockTypeWord {
|
||||
if str.String() != "" {
|
||||
str.WriteRune(' ')
|
||||
}
|
||||
str.WriteString(*blockMap[uid].Text)
|
||||
}
|
||||
|
||||
if blockMap[uid].BlockType != types.BlockTypeSelectionElement {
|
||||
continue
|
||||
}
|
||||
|
||||
if isSelected == nil || !*isSelected {
|
||||
selected := true
|
||||
if blockMap[uid].SelectionStatus == types.SelectionStatusNotSelected {
|
||||
selected = false
|
||||
}
|
||||
isSelected = &selected
|
||||
}
|
||||
}
|
||||
|
||||
selected := true
|
||||
if isSelected != nil {
|
||||
selected = *isSelected
|
||||
}
|
||||
|
||||
return str.String(), selected
|
||||
}
|
||||
|
||||
func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, isSelected bool, params *tableParams) {
|
||||
if params.maxRow < rowIndex {
|
||||
params.maxRow = rowIndex
|
||||
}
|
||||
if params.maxColumn < columnIndex {
|
||||
params.maxColumn = columnIndex
|
||||
}
|
||||
params.table = append(params.table, cell{
|
||||
column: columnIndex,
|
||||
row: rowIndex,
|
||||
value: content,
|
||||
isSelected: isSelected,
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user