query-orchestration/api/docTextRunner/runner_test.go

package doctextrunner_test

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"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"
	objectstoremock "queryorchestration/mocks/objectstore"
	queuemock "queryorchestration/mocks/queue"
	textractmock "queryorchestration/mocks/textract"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sqs"
	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
	"github.com/google/uuid"
	"github.com/pashagolub/pgxmock/v3"
	"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 TestDocCleanRunner(t *testing.T) {
	ctx := context.Background()
	pool, err := pgxmock.NewPool()
	require.NoError(t, err)

	cfg := &DocTextConfig{}
	cfg.DBPool = pool
	cfg.DBQueries = repository.New(pool)
	mockStore := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockStore
	mockTextract := textractmock.NewMockTextractClient(t)
	cfg.TextractClient = mockTextract
	mockQueue := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockQueue
	cfg.QuerySyncURL = "ex"

	runner := doctextrunner.New(&doctextrunner.Services{
		Text: documenttext.New(cfg),
	})
	assert.NotNil(t, runner)
	doc := doctextrunner.Body{
		DocumentID: uuid.New(),
	}

	cleanId := uuid.New()
	bucket := "bucket"
	hash := `"a06d5e0e795adeb6d0b3732330dbcc15"`
	textId := uuid.New()
	key := objectstore.BucketKey{
		ClientID:  "aa",
		Location:  objectstore.Import,
		CreatedAt: time.Now().UTC(),
		EntityID:  cleanId,
	}
	keyStr := key.String()
	filename := "../../assets/sampleGeneration/generated/helloWorld.gen"
	file, err := os.Open(filename)
	require.NoError(t, err)
	defer file.Close()

	decoder := json.NewDecoder(file)
	var textractBaseOut map[string][]awstextract.AnalyzeDocumentOutput
	err = decoder.Decode(&textractBaseOut)
	require.NoError(t, err)

	mimetype := repository.NullCleanmimetype{
		Valid:         true,
		Cleanmimetype: repository.CleanmimetypeApplicationPdf,
	}
	pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(doc.DocumentID).WillReturnRows(
		pgxmock.NewRows([]string{"isextracted"}).
			AddRow(false),
	)
	pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(doc.DocumentID).
		WillReturnRows(
			pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}).
				AddRow(cleanId, doc.DocumentID, &bucket, &keyStr, int64(123), &hash, mimetype, nil, key.ClientID),
		)
	bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
	mockStore.EXPECT().
		GetObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
				return *in.Bucket == bucket && *in.Key == keyStr
			}),
			mock.Anything,
		).
		Return(&s3.GetObjectOutput{
			Body: bodyMsg,
		}, nil)
	sent := 0
	mockTextract.EXPECT().
		AnalyzeDocument(
			mock.Anything,
			mock.MatchedBy(func(in *awstextract.AnalyzeDocumentInput) bool {
				if sent == 0 {
					sent++
					return true
				}
				return false
			}),
			mock.Anything,
		).
		Return(&textractBaseOut["base"][0], nil).Once()

	pool.ExpectBegin()
	pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows(
		pgxmock.NewRows([]string{"id"}).
			AddRow(textId),
	)
	pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
		WillReturnResult(pgxmock.NewResult("", 1))
	pool.ExpectCommit()

	mockQueue.EXPECT().
		SendMessage(
			mock.Anything,
			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
				return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.DocumentID.String())
			}),
			mock.Anything,
		).
		Return(&sqs.SendMessageOutput{}, nil)

	assert.True(t, runner.Process(ctx, doc))
}

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`