query-orchestration/internal/document/text/pageText_test.go

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 := "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`