query-orchestration/api/docCleanRunner/runner_test.go

package doccleanrunner_test

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

	doccleanrunner "queryorchestration/api/docCleanRunner"
	"queryorchestration/internal/database/repository"
	documentclean "queryorchestration/internal/document/clean"
	"queryorchestration/internal/server/runner"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue/documenttext"
	"queryorchestration/internal/test"
	objectstoremock "queryorchestration/mocks/objectstore"

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

type DocCleanConfig struct {
	runner.BaseConfig[doccleanrunner.Body]
	documenttext.DocTextConfig
	objectstore.ObjectStoreConfig
}

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

	test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint)
	cfg.DocumentTextURL = test.CreateQueue(t, cfg, test.DocTextRunnerName)

	mockStore := objectstoremock.NewMockS3Client(t)
	cfg.StoreClient = mockStore

	runner := doccleanrunner.New(&doccleanrunner.Services{
		Clean: documentclean.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)

	mimeType := "application/pdf"
	mockStore.EXPECT().
		HeadObject(
			mock.Anything,
			mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
				return *in.Bucket == "bucket" && *in.Key == key.String()
			}),
			mock.Anything,
		).
		Return(&s3.HeadObjectOutput{
			ContentType: &mimeType,
		}, nil)

	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)

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

	test.AssertMessageBody(t, cfg, cfg.GetDocumentTextURL(), 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`