package documentclean import ( "context" "fmt" "io" "strings" "testing" "queryorchestration/internal/database" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" cleanversion "queryorchestration/internal/document/clean/version" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documenttext" 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" ) type DocCleanConfig struct { serviceconfig.BaseConfig documenttext.DocTextConfig objectstore.ObjectStoreConfig } func TestCreate(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() require.NoError(t, err) mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocCleanConfig{} cfg.QueueClient = mockSQS cfg.DocumentTextURL = "/i/am/here" cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockS3 := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockS3 svc := Service{ cfg: cfg, svc: &Services{ Version: cleanversion.New(cfg), }, } doc := document.Document{ ID: uuid.New(), JobID: uuid.New(), Hash: "example_hash", } inloc := document.Location{ Bucket: "bucket_name", Key: "/i/am/here", } dbid := database.MustToDBUUID(doc.ID) pool.ExpectQuery("name: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"isclean"}). AddRow(false), ) pool.ExpectQuery("name: GetDocument :one").WithArgs(dbid). WillReturnRows( pgxmock.NewRows([]string{"id", "jobId", "hash"}). AddRow(dbid, database.MustToDBUUID(doc.JobID), doc.Hash), ) pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows( pgxmock.NewRows([]string{"documentId", "bucket", "key"}). AddRow(dbid, inloc.Bucket, inloc.Key), ) 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", "mimetype", "fail"}), ) cleanid := database.MustToDBUUID(uuid.New()) pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &inloc.Bucket, &inloc.Key, dbmimetype, repository.NullCleanfailtype{}). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(cleanid), ) pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, int32(1)). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.DocumentTextURL && *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 == inloc.Bucket && *in.Key == inloc.Key }), 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 == inloc.Bucket && *in.Key == inloc.Key }), mock.Anything, ). Return(&s3.GetObjectOutput{ Body: body, }, nil) err = svc.Clean(ctx, doc.ID) assert.NoError(t, err) } func TestInformClean(t *testing.T) { ctx := context.Background() mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocCleanConfig{} cfg.QueueClient = mockSQS cfg.DocumentTextURL = "/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.DocumentTextURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) err := svc.informClean(ctx, id) assert.NoError(t, err) }