query-orchestration/api/docTextRunner/runner_test.go

package doctextrunner_test

import (
	"fmt"
	"io"
	"regexp"
	"strings"
	"testing"
	"time"

	doctextrunner "queryorchestration/api/docTextRunner"
	"queryorchestration/internal/database/repository"
	documenttext "queryorchestration/internal/document/text"
	"queryorchestration/internal/server/runner"
	"queryorchestration/internal/serviceconfig/aws"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue/querysync"
	"queryorchestration/internal/serviceconfig/textract"
	"queryorchestration/internal/test"
	objectstoremock "queryorchestration/mocks/objectstore"
	textractmock "queryorchestration/mocks/textract"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

type DocTextConfig struct {
	runner.BaseConfig[doctextrunner.Body]
	aws.AWSConfig
	querysync.QuerySyncConfig
	textract.TextractConfig
	objectstore.ObjectStoreConfig
}

func TestDocTextRunner(t *testing.T) {
	cfg := &DocTextConfig{}
	test.CreateDB(t, cfg)
	acfg := test.CreateAWSContainer(t, cfg)

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.QuerySyncURL = test.CreateQueue(t, cfg, test.QuerySyncRunnerName)

	mockStore := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockStore
	mockTextract := textractmock.NewMockTextractClient(t)
	cfg.TextractClient = mockTextract

	runner := doctextrunner.New(&doctextrunner.Services{
		Text: documenttext.New(cfg),
	})

	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
		Clientid: "clientid",
		Name:     "client_name",
	})
	require.NoError(t, err)
	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
		Clientid: "clientid",
		Hash:     "hash",
	})
	require.NoError(t, err)
	part := uint16(1)
	filetype := "pdf"
	key := objectstore.BucketKey{
		CreatedAt: time.Now().UTC(),
		ClientID:  "clientid",
		EntityID:  uuid.New(),
		Location:  objectstore.Import,
		Part:      &part,
		FileType:  &filetype,
	}
	err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
		Documentid: docId,
		Bucket:     "bucket",
		Key:        key.String(),
	})
	require.NoError(t, err)
	bucket := "bucket"
	hash := "hash"
	keyStr := key.String()
	cleanId, err := cfg.GetDBQueries().AddDocumentClean(t.Context(), &repository.AddDocumentCleanParams{
		Documentid: docId,
		Bucket:     &bucket,
		Hash:       &hash,
		Key:        &keyStr,
		Mimetype: repository.NullCleanmimetype{
			Valid:         true,
			Cleanmimetype: repository.CleanmimetypeApplicationPdf,
		},
	})
	require.NoError(t, err)
	err = cfg.GetDBQueries().AddDocumentCleanEntry(t.Context(), &repository.AddDocumentCleanEntryParams{
		Cleanid: cleanId,
		Version: 1,
	})
	require.NoError(t, err)

	bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
	mockStore.EXPECT().
		GetObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
				return *in.Bucket == "bucket" && *in.Key == key.String()
			}),
			mock.Anything,
		).
		Return(&s3.GetObjectOutput{
			Body: bodyMsg,
		}, nil)
	filename := "../../assets/textract/helloWorld.gen"
	textractBaseOut := test.GetTextractFileResponse(t, filename)
	mockTextract.EXPECT().
		AnalyzeDocument(
			mock.Anything,
			mock.MatchedBy(func(in *awstextract.AnalyzeDocumentInput) bool {
				return true
			}),
			mock.Anything,
		).
		Return(textractBaseOut[0]["base"], nil).Once()
	mockStore.EXPECT().
		PutObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.PutObjectInput) bool {
				return *in.Bucket == "bucket"
			}),
			mock.Anything,
		).
		Return(&s3.PutObjectOutput{}, nil)

	doc := doctextrunner.Body{
		DocumentID: docId,
	}
	assert.True(t, runner.Process(t.Context(), doc))

	test.AssertMessageBody(t, cfg, cfg.GetQuerySyncURL(), regexp.MustCompile(fmt.Sprintf("{\"id\":\"%s\"}", docId.String())))
}

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`