package documentclean import ( "fmt" "io" "strings" "testing" "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/aws" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documenttexttrigger" 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/mock" "github.com/stretchr/testify/require" ) type DocCleanConfig struct { aws.AWSConfig serviceconfig.BaseConfig documenttexttrigger.DocTextTriggerConfig objectstore.ObjectStoreConfig } func TestCreate(t *testing.T) { ctx := t.Context() pool, err := pgxmock.NewPool() require.NoError(t, err) mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocCleanConfig{} cfg.QueueClient = mockSQS cfg.DocumentTextTriggerURL = "/i/am/here" cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockS3 := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockS3 svc := Service{ cfg: cfg, } doc := document.DocumentSummary{ ID: uuid.New(), ClientID: "yuyu", Hash: "example_hash", } bucket := "example_bucket" key := objectstore.BucketKey{ ClientID: doc.ClientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), Location: objectstore.Import, } keyStr := key.String() dbid := doc.ID pool.ExpectQuery("name: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"isclean"}). AddRow(false), ) pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "hash"}). AddRow(dbid, doc.ClientID, doc.Hash), ) pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key"}). AddRow(dbid, bucket, keyStr), ) mimeType := "application/pdf" dbmimetype := repository.NullCleanmimetype{ Valid: true, Cleanmimetype: repository.Cleanmimetype(mimeType), } pool.ExpectBegin() pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}), ) cleanid := uuid.New() pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &bucket, &keyStr, &doc.Hash, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.DocumentTextTriggerURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) mockS3.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) body := io.NopCloser(strings.NewReader(pdfHelloWorld)) mockS3.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: body, }, nil) err = svc.Clean(ctx, doc.ID) require.NoError(t, err) } func TestInformClean(t *testing.T) { ctx := t.Context() mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocCleanConfig{} cfg.QueueClient = mockSQS cfg.DocumentTextTriggerURL = "/i/am/here" svc := Service{ cfg: cfg, } id := uuid.New() mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.DocumentTextTriggerURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err := svc.informClean(ctx, id) require.NoError(t, err) }