Michael McGuinness
d54c823d77
Stabilise tests and CNC_1 Sample * testspeedup * notyet * replacebits * noflake
460 lines
9.9 KiB
Go
460 lines
9.9 KiB
Go
package documenttext
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"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)
|
|
|
|
root, err := os.OpenRoot("../../../assets")
|
|
require.NoError(t, err)
|
|
textractBaseOut := getTextractFileResponse(t, "sampleGeneration/generated/helloWorld.gen", root)
|
|
|
|
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()
|
|
|
|
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, txtFile, close := getFiles(t, "sampleGeneration", "helloWorld")
|
|
defer close()
|
|
|
|
textExpectations(t, mockTextract, textractBaseOut)
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
assertReaders(t, txtFile, text)
|
|
})
|
|
|
|
t.Run("sampleOne - cnc_01-06", func(t *testing.T) {
|
|
t.Parallel()
|
|
cfg := DocTextConfig{}
|
|
svc := Service{
|
|
cfg: &cfg,
|
|
}
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
|
|
textractBaseOut, pdf, txtFile, close := getFiles(t, "sampleOne", "cnc_01-06")
|
|
defer close()
|
|
|
|
textExpectations(t, mockTextract, textractBaseOut)
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
|
|
assertReaders(t, txtFile, text)
|
|
})
|
|
t.Run("sampleOne - hn_0109", func(t *testing.T) {
|
|
t.Parallel()
|
|
cfg := DocTextConfig{}
|
|
svc := Service{
|
|
cfg: &cfg,
|
|
}
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
|
|
textractBaseOut, pdf, txtFile, close := getFiles(t, "sampleOne", "hn_0109")
|
|
defer close()
|
|
|
|
textExpectations(t, mockTextract, textractBaseOut)
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
|
|
assertReaders(t, txtFile, text)
|
|
})
|
|
t.Run("sampleOne - hn23-70", func(t *testing.T) {
|
|
t.Parallel()
|
|
cfg := DocTextConfig{}
|
|
svc := Service{
|
|
cfg: &cfg,
|
|
}
|
|
mockTextract := textractmock.NewMockTextractClient(t)
|
|
cfg.TextractClient = mockTextract
|
|
|
|
textractBaseOut, pdf, txtFile, close := getFiles(t, "sampleOne", "hn_23-70")
|
|
defer close()
|
|
|
|
textExpectations(t, mockTextract, textractBaseOut)
|
|
|
|
text, err := svc.getDocumentText(ctx, pdf)
|
|
require.NoError(t, err)
|
|
|
|
assertReaders(t, txtFile, text)
|
|
})
|
|
// t.Run("sampleOne - chc_1", func(t *testing.T) {
|
|
// t.Parallel()
|
|
// cfg := DocTextConfig{}
|
|
// svc := Service{
|
|
// cfg: &cfg,
|
|
// }
|
|
// mockTextract := textractmock.NewMockTextractClient(t)
|
|
// cfg.TextractClient = mockTextract
|
|
|
|
// textractBaseOut, pdf, txtFile, close := getFiles(t, "sampleOne", "chc_1")
|
|
// defer close()
|
|
|
|
// textExpectations(t, mockTextract, textractBaseOut)
|
|
|
|
// text, err := svc.getDocumentText(ctx, pdf)
|
|
// require.NoError(t, err)
|
|
|
|
// assertReaders(t, txtFile, text)
|
|
// })
|
|
}
|
|
|
|
func getFiles(t testing.TB, sample string, filename string) (map[string][]textract.AnalyzeDocumentOutput, *documenttypes.PDF, *os.File, func()) {
|
|
root, err := os.OpenRoot("../../../assets")
|
|
require.NoError(t, err)
|
|
root, err = root.OpenRoot(sample)
|
|
require.NoError(t, err)
|
|
|
|
textractFile, err := root.Open(fmt.Sprintf("generated/%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("%s.pdf", filename))
|
|
require.NoError(t, err)
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(pdfFile)
|
|
require.NoError(t, err)
|
|
|
|
txtFile, err := root.Open(fmt.Sprintf("%s.txt", filename))
|
|
require.NoError(t, err)
|
|
|
|
return out, pdf, txtFile, func() {
|
|
pdfFile.Close()
|
|
txtFile.Close()
|
|
}
|
|
}
|
|
|
|
func textExpectations(t testing.TB, mockTextract *textractmock.MockTextractClient, textractBaseOut map[string][]textract.AnalyzeDocumentOutput) {
|
|
sent := 0
|
|
tableIndex := 0
|
|
for i := range len(textractBaseOut["base"]) {
|
|
curr := i + tableIndex
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == curr {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["base"][i], nil).Once()
|
|
|
|
hasTable := false
|
|
for _, t := range textractBaseOut["base"][i].Blocks {
|
|
if t.BlockType == types.BlockTypeLayoutTable {
|
|
hasTable = true
|
|
break
|
|
}
|
|
}
|
|
|
|
if hasTable {
|
|
curr := i + tableIndex + 1
|
|
mockTextract.EXPECT().
|
|
AnalyzeDocument(
|
|
mock.Anything,
|
|
mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
|
|
if sent == curr && in.FeatureTypes[0] == types.FeatureTypeTables {
|
|
sent++
|
|
return true
|
|
}
|
|
return false
|
|
}),
|
|
mock.Anything,
|
|
).
|
|
Return(&textractBaseOut["table"][tableIndex], nil).Once()
|
|
tableIndex++
|
|
}
|
|
}
|
|
}
|
|
|
|
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.TB, filename string, root *os.Root) map[string][]textract.AnalyzeDocumentOutput {
|
|
file, err := root.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 assertReaders(t testing.TB, 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)
|
|
}
|
|
}
|
|
|
|
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`
|