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

package documenttext

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"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/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestGetDocumentText(t *testing.T) {
	ctx := t.Context()

	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, "helloWorld")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)
		assertReaders(t, txtFile, text)
	})
	t.Run("Error Getting Textract Response", func(t *testing.T) {
		t.Parallel()
		cfg := DocTextConfig{}
		svc := Service{
			cfg: &cfg,
		}
		mockTextract := textractmock.NewMockTextractClient(t)
		cfg.TextractClient = mockTextract

		_, pdf, _, close := getFiles(t, "helloWorld")
		defer close()

		mockTextract.EXPECT().
			AnalyzeDocument(
				mock.Anything,
				mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
					return true
				}),
				mock.Anything,
			).
			Return(nil, errors.New("no textract response"))

		_, err := svc.GetDocumentText(ctx, pdf)
		require.EqualError(t, err, "no textract response")
	})
	t.Run("merged_cell_table", 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, "merged_cell_table")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("table_check", 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, "table_check")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("table_select", 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, "table_select")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("multi_column", 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, "multi_column")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("table_check_row", 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, "table_check_row")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("table_aligned", 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, "table_aligned")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})

	t.Run("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, "cnc_01-06")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("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, "hn_0109")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("hn_23-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, "hn_23-70")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
	t.Run("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, "chc_1")
		defer close()

		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)

		text, err := svc.GetDocumentText(ctx, pdf)
		require.NoError(t, err)

		assertReaders(t, txtFile, text)
	})
}

func getFiles(t testing.TB, filename string) ([]map[string]*textract.AnalyzeDocumentOutput, *documenttypes.PDF, *os.File, func()) {
	root, err := os.OpenRoot("../../../assets")
	require.NoError(t, err)

	textractFile, err := root.Open(fmt.Sprintf("textract/%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("original/%s.pdf", filename))
	require.NoError(t, err)

	pdf, err := documenttypes.NewPDFFromReader(pdfFile)
	require.NoError(t, err)

	txtFile, err := root.Open(fmt.Sprintf("text/%s.txt", filename))
	require.NoError(t, err)

	return out, pdf, txtFile, func() {
		pdfFile.Close()
		txtFile.Close()
	}
}

func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractmock.MockTextractClient, pdf documenttypes.File, textractBaseOut []map[string]*textract.AnalyzeDocumentOutput) {
	for i, pageResult := range textractBaseOut {
		page, err := pdf.GetPageAsPNG(ctx, i)
		require.NoError(t, err)

		base := pageResult["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)

		table := pageResult["full_features"]
		if table != nil {
			mockTextract.EXPECT().
				AnalyzeDocument(
					mock.Anything,
					mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
						found := false
						for _, feature := range in.FeatureTypes {
							if feature == types.FeatureTypeForms || feature == types.FeatureTypeTables {
								found = true
								break
							}
						}
						return string(in.Document.Bytes) == string(page) && found
					}),
					mock.Anything,
				).
				Return(table, nil)
		}
	}
}

func assertReaderAndString(t testing.TB, expected io.Reader, actual string) {
	expectedBytes, err := io.ReadAll(expected)
	require.NoError(t, err)
	expectedStr := string(expectedBytes)

	if !assert.Equal(t, expectedStr, actual) {
		diff := cmp.Diff(expectedStr, actual)
		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
	}
}

func assertReaders(t testing.TB, expected io.Reader, actual io.Reader) {
	actualBytes, err := io.ReadAll(actual)
	require.NoError(t, err)
	actualStr := string(actualBytes)

	assertReaderAndString(t, expected, actualStr)
}