Jay Brown
330 lines
6.6 KiB
Go
330 lines
6.6 KiB
Go
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`
|