query-orchestration/internal/document/text/getPages_test.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`