package documenttext import ( "context" "io" "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 TestGetBasePage(t *testing.T) { ctx := context.Background() cfg := DocTextConfig{} svc := Service{ cfg: &cfg, } mockTextract := textractmock.NewMockTextractClient(t) cfg.TextractClient = mockTextract textractBaseOut, pdf, _, close := getFiles(t, "helloWorld") defer close() page, err := pdf.GetPageAsPNG(ctx, 0) require.NoError(t, err) base := textractBaseOut[0]["base"] mockTextract.EXPECT(). AnalyzeDocument( mock.Anything, mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool { return string(in.Document.Bytes) == string(page) && 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 := "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)) }) } 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`