Michael McGuinness
fee71e7740
Feature/postprocessing * tests * passtest * fixshorttests * mosttests * improvingbasedockerfile * testspeeds * testing * host * canparallel * clean * passfullsuite * singlepagemax * test * findfeatures * findstables * tbls * tablestoo * tablestoo * lateraltests * tableloc * cleanup * inlinetable * childids * cleanup * tests
571 lines
12 KiB
Go
571 lines
12 KiB
Go
package documenttext
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"math"
|
|
"os"
|
|
"strings"
|
|
"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/google/go-cmp/cmp"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGetBasePage(t *testing.T) {
|
|
ctx := context.Background()
|
|
cfg := DocTextConfig{}
|
|
svc := Service{
|
|
cfg: &cfg,
|
|
}
|
|
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
sent := 0
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(strings.NewReader(pdfHelloWorld))
|
|
require.NoError(t, err)
|
|
|
|
textractFile := "../../../assets/sampleGeneration/generated/helloWorld.gen"
|
|
textractBaseOut := getTextractFileResponse(t, textractFile)
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 0 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][0], nil).Once()
|
|
|
|
text, err := svc.GetBasePage(ctx, pdf, 0)
|
|
require.NoError(t, err)
|
|
assert.EqualExportedValues(t, textractBaseOut["base"][0], text)
|
|
}
|
|
|
|
func TestGetDocumentText(t *testing.T) {
|
|
ctx := context.Background()
|
|
cfg := DocTextConfig{}
|
|
svc := Service{
|
|
cfg: &cfg,
|
|
}
|
|
|
|
t.Run("Hello World", func(t *testing.T) {
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
sent := 0
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(strings.NewReader(pdfHelloWorld))
|
|
require.NoError(t, err)
|
|
|
|
textractFile := "../../../assets/sampleGeneration/generated/helloWorld.gen"
|
|
textractBaseOut := getTextractFileResponse(t, textractFile)
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 0 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][0], nil).Once()
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
expected := `Start of Page No. = 1
|
|
|
|
Hello World!
|
|
|
|
|
|
`
|
|
assertStringToReader(t, expected, text)
|
|
})
|
|
|
|
t.Run("sampleOne - cnc_01-06", func(t *testing.T) {
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
sent := 0
|
|
|
|
textractFile := "../../../assets/sampleOne/generated/cnc_01-06.gen"
|
|
textractBaseOut := getTextractFileResponse(t, textractFile)
|
|
|
|
pdfFile, err := os.Open("../../../assets/sampleOne/cnc_01-06.pdf")
|
|
require.NoError(t, err)
|
|
defer pdfFile.Close()
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(pdfFile)
|
|
require.NoError(t, err)
|
|
|
|
txtFile, err := os.Open("../../../assets/sampleOne/cnc_01-06.txt")
|
|
require.NoError(t, err)
|
|
defer txtFile.Close()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 0 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][0], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 1 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][1], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 2 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][2], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 3 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][3], nil).Once()
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
|
|
assertReaders(t, txtFile, text)
|
|
})
|
|
t.Run("sampleOne - hn_0109", func(t *testing.T) {
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
sent := 0
|
|
|
|
textractFile := "../../../assets/sampleOne/generated/hn_0109.gen"
|
|
textractBaseOut := getTextractFileResponse(t, textractFile)
|
|
|
|
pdfFile, err := os.Open("../../../assets/sampleOne/hn_0109.pdf")
|
|
require.NoError(t, err)
|
|
defer pdfFile.Close()
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(pdfFile)
|
|
require.NoError(t, err)
|
|
|
|
txtFile, err := os.Open("../../../assets/sampleOne/hn_0109.txt")
|
|
require.NoError(t, err)
|
|
defer txtFile.Close()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 0 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][0], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 1 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][1], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 2 && in.FeatureTypes[0] == types.FeatureTypeTables {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["table"][0], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 3 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][2], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 4 && in.FeatureTypes[0] == types.FeatureTypeTables {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["table"][1], nil).Once()
|
|
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == 5 {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][3], nil).Once()
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
|
|
assertReaders(t, txtFile, 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 := "blockId.String()"
|
|
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))
|
|
})
|
|
}
|
|
|
|
func getTextractFileResponse(t *testing.T, filename string) map[string][]textract.AnalyzeDocumentOutput {
|
|
file, err := os.Open(filename)
|
|
require.NoError(t, err)
|
|
defer file.Close()
|
|
|
|
decoder := json.NewDecoder(file)
|
|
var out map[string][]textract.AnalyzeDocumentOutput
|
|
err = decoder.Decode(&out)
|
|
require.NoError(t, err)
|
|
|
|
return out
|
|
}
|
|
|
|
func assertStringToReader(t *testing.T, expected string, actual io.Reader) {
|
|
actualBytes, err := io.ReadAll(actual)
|
|
require.NoError(t, err)
|
|
actualStr := string(actualBytes)
|
|
if !assert.Equal(t, expected, actualStr) {
|
|
diff := cmp.Diff(expected, actualStr)
|
|
t.Logf("Detailed diff (-expected +actual):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
func assertReaders(t *testing.T, expected io.Reader, actual io.Reader) {
|
|
actualBytes, err := io.ReadAll(actual)
|
|
require.NoError(t, err)
|
|
actualStr := string(actualBytes)
|
|
|
|
expectedBytes, err := io.ReadAll(expected)
|
|
require.NoError(t, err)
|
|
expectedStr := string(expectedBytes)
|
|
|
|
if !assert.Equal(t, expectedStr, actualStr) {
|
|
diff := cmp.Diff(expectedStr, actualStr)
|
|
t.Logf("Detailed diff (-expected +actual):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
func levenshteinDistance(s1, s2 string) int {
|
|
rows, cols := len(s1)+1, len(s2)+1
|
|
dist := make([][]int, rows)
|
|
for i := range dist {
|
|
dist[i] = make([]int, cols)
|
|
}
|
|
|
|
for i := 0; i < rows; i++ {
|
|
dist[i][0] = i
|
|
}
|
|
for j := 0; j < cols; j++ {
|
|
dist[0][j] = j
|
|
}
|
|
|
|
for i := 1; i < rows; i++ {
|
|
for j := 1; j < cols; j++ {
|
|
cost := 1
|
|
if s1[i-1] == s2[j-1] {
|
|
cost = 0
|
|
}
|
|
dist[i][j] = min(
|
|
dist[i-1][j]+1,
|
|
dist[i][j-1]+1,
|
|
dist[i-1][j-1]+cost,
|
|
)
|
|
}
|
|
}
|
|
|
|
return dist[rows-1][cols-1]
|
|
}
|
|
|
|
func calculateSimilarityPercentage(s1, s2 string) float64 {
|
|
distance := levenshteinDistance(s1, s2)
|
|
maxLen := math.Max(float64(len(s1)), float64(len(s2)))
|
|
|
|
if maxLen == 0 {
|
|
return 100.0
|
|
}
|
|
|
|
return 100.0 * (1.0 - float64(distance)/maxLen)
|
|
}
|
|
|
|
func AreSimilar(s1, s2 string, threshold float64) bool {
|
|
return calculateSimilarityPercentage(s1, s2) >= threshold
|
|
}
|
|
|
|
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`
|