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

package documenttext

import (
	"context"
	"database/sql"
	"fmt"
	"io"
	"strings"
	"testing"
	"time"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig/objectstore"
	objectstoremock "queryorchestration/mocks/objectstore"
	queuemock "queryorchestration/mocks/queue"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sqs"
	"github.com/google/uuid"
	"github.com/pashagolub/pgxmock/v3"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

func TestProcess(t *testing.T) {
	ctx := context.Background()
	pool, err := pgxmock.NewPool()
	require.NoError(t, err)

	cfg := &DocTextConfig{}
	cfg.DBPool = pool
	cfg.DBQueries = repository.New(pool)
	cfg.QuerySyncURL = "querysync"
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockS3 := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockS3

	svc := Service{
		cfg: cfg,
	}

	triggerId := uuid.MustParse("6ac21b2b-f36c-4ae7-a815-307d4a9bfe4d")
	docId := uuid.MustParse("6ac21b2b-f36c-4ae7-a815-307d4a9bfe4a")
	textractId := "hello"
	textId := uuid.New()
	cleanId := uuid.New()
	hash := `"5d41402abc4b2a76b9719d911017c592"`

	pool.ExpectQuery("name: GetDocumentTextTrigger :one").WithArgs(triggerId).WillReturnRows(
		pgxmock.NewRows([]string{"id", "cleanId", "version", "textractJobId", "documentId"}).
			AddRow(triggerId, cleanId, int64(123), &textractId, docId),
	)
	pool.ExpectBegin()
	pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows(
		pgxmock.NewRows([]string{"id"}).
			AddRow(textId),
	)
	pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, triggerId, pgxmock.AnyArg()).
		WillReturnResult(pgxmock.NewResult("", 1))
	pool.ExpectCommit()

	body := io.NopCloser(strings.NewReader("hello"))
	mockS3.EXPECT().
		GetObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
				return *in.IfMatch == "hash"
			}),
			mock.Anything,
		).
		Return(&s3.GetObjectOutput{
			Body: body,
		}, nil)

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

	key := objectstore.BucketKey{
		ClientID: "my_client",
		Location: objectstore.TextTextract,
		EntityID: triggerId,
	}
	err = svc.Process(ctx, &ProcessParams{
		Bucket: "bucket",
		Key:    key,
		Hash:   "hash",
	})
	assert.NoError(t, err)
}

func TestStoreTrigger(t *testing.T) {
	ctx := context.Background()
	pool, err := pgxmock.NewPool()
	require.NoError(t, err)

	cfg := &DocTextConfig{}
	cfg.DBPool = pool
	cfg.DBQueries = repository.New(pool)
	cfg.QuerySyncURL = "querysync"
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockS3 := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockS3

	svc := Service{
		cfg: cfg,
	}

	cleanId := uuid.New()
	textId := uuid.New()
	triggerId := uuid.New()
	clientId := "hello"
	bucket := "bucket"

	t.Run("exists", func(t *testing.T) {
		hash := `"5d41402abc4b2a76b9719d911017c592"`
		pool.ExpectBegin()
		pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows(
			pgxmock.NewRows([]string{"id"}).
				AddRow(textId),
		)
		pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, triggerId, pgxmock.AnyArg()).
			WillReturnResult(pgxmock.NewResult("", 1))
		pool.ExpectCommit()

		text := strings.NewReader("hello")
		err = svc.storeProcess(ctx, text, &repository.GetDocumentTextTriggerRow{
			ID:      triggerId,
			Cleanid: cleanId,
		}, &ProcessParams{
			Bucket: bucket,
			Key:    objectstore.BucketKey{},
		})
		require.NoError(t, err)
	})
	t.Run("not exists", func(t *testing.T) {
		part := uint16(0)
		key := objectstore.BucketKey{
			ClientID:  clientId,
			Location:  objectstore.TextOut,
			CreatedAt: time.Now().UTC(),
			EntityID:  triggerId,
			Part:      &part,
		}
		hash := `"09644372e99020106946045c6fd2d70b"`
		pool.ExpectBegin()
		pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).
			WillReturnError(sql.ErrNoRows)
		pool.ExpectQuery("name: GetTextOutCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()).
			WillReturnRows(
				pgxmock.NewRows([]string{"part", "count"}).
					AddRow(0, 2),
			)
		pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash, pgxmock.AnyArg(), uint16(0)).
			WillReturnRows(
				pgxmock.NewRows([]string{"id"}).
					AddRow(textId),
			)

		mockS3.EXPECT().
			PutObject(
				mock.Anything,
				mock.MatchedBy(func(in *s3.PutObjectInput) bool {
					bod, err := io.ReadAll(in.Body)
					require.NoError(t, err)
					return string(bod) == "byebye" && *in.Key == key.String() && *in.Bucket == bucket
				}),
				mock.Anything,
			).
			Return(&s3.PutObjectOutput{}, nil)

		pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, triggerId, pgxmock.AnyArg()).
			WillReturnResult(pgxmock.NewResult("", 1))
		pool.ExpectCommit()

		text := strings.NewReader("byebye")
		err = svc.storeProcess(ctx, text, &repository.GetDocumentTextTriggerRow{
			ID:      triggerId,
			Cleanid: cleanId,
		}, &ProcessParams{
			Key:    key,
			Bucket: bucket,
		})
		require.NoError(t, err)
	})
}

func TestDuplicateReader(t *testing.T) {
	original := strings.NewReader("hello")
	readerOne, readerTwo, err := duplicateReader(original)
	require.NoError(t, err)
	contentOne, err := io.ReadAll(readerOne)
	require.NoError(t, err)
	assert.Equal(t, []byte("hello"), contentOne)
	contentTwo, err := io.ReadAll(readerTwo)
	require.NoError(t, err)
	assert.Equal(t, []byte("hello"), contentTwo)
}

func TestProcessTrigger(t *testing.T) {
	ctx := context.Background()
	pool, err := pgxmock.NewPool()
	require.NoError(t, err)

	cfg := &DocTextConfig{}
	cfg.DBPool = pool
	cfg.DBQueries = repository.New(pool)
	cfg.QuerySyncURL = "querysync"
	mockSQS := queuemock.NewMockSQSClient(t)
	cfg.QueueClient = mockSQS
	mockS3 := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockS3

	svc := Service{
		cfg: cfg,
	}

	cleanId := uuid.New()
	textId := uuid.New()
	triggerId := uuid.New()
	clientId := "hello"
	bucket := "bucket"
	key := objectstore.BucketKey{
		ClientID:  clientId,
		Location:  objectstore.TextOut,
		EntityID:  triggerId,
		CreatedAt: time.Now().UTC(),
	}
	hash := `"5d41402abc4b2a76b9719d911017c592"`

	pool.ExpectBegin()
	pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).
		WillReturnError(sql.ErrNoRows)
	pool.ExpectQuery("name: GetTextOutCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()).
		WillReturnRows(
			pgxmock.NewRows([]string{"part", "count"}).
				AddRow(0, 2),
		)
	pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash, pgxmock.AnyArg(), uint16(0)).
		WillReturnRows(
			pgxmock.NewRows([]string{"id"}).
				AddRow(textId),
		)

	body := io.NopCloser(strings.NewReader("hello"))
	mockS3.EXPECT().
		GetObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
				return *in.Key == key.String() && *in.Bucket == bucket && *in.IfMatch == hash
			}),
			mock.Anything,
		).
		Return(&s3.GetObjectOutput{
			Body: body,
		}, nil)

	mockS3.EXPECT().
		PutObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.PutObjectInput) bool {
				bod, err := io.ReadAll(in.Body)
				require.NoError(t, err)
				return string(bod) == "hello" && *in.Key == key.String() && *in.Bucket == bucket
			}),
			mock.Anything,
		).
		Return(&s3.PutObjectOutput{}, nil)

	pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, triggerId, pgxmock.AnyArg()).
		WillReturnResult(pgxmock.NewResult("", 1))
	pool.ExpectCommit()

	err = svc.processTrigger(ctx, &repository.GetDocumentTextTriggerRow{
		ID:      triggerId,
		Cleanid: cleanId,
	}, &ProcessParams{
		Bucket: bucket,
		Key:    key,
		Hash:   hash,
	})
	require.NoError(t, err)
}