package documentinit import ( "context" "errors" "fmt" "queryorchestration/internal/client" "queryorchestration/internal/database" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" "queryorchestration/internal/job" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documentclean" objectstoremock "queryorchestration/mocks/objectstore" queuemock "queryorchestration/mocks/queue" "testing" "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" ) type DocInitConfig struct { serviceconfig.BaseConfig documentclean.DocCleanConfig objectstore.ObjectStoreConfig } func TestCreate(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) j := job.Job{ ID: uuid.New(), ClientID: uuid.New(), } doc := document.Document{ ID: uuid.New(), JobID: j.ID, Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectQuery("name: GetJob :one").WithArgs(database.MustToDBUUID(j.ID)).WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "canSync"}). AddRow(database.MustToDBUUID(j.ID), database.MustToDBUUID(j.ClientID), j.CanSync), ) pool.ExpectQuery("-- name: GetClient :one").WithArgs(database.MustToDBUUID(j.ClientID)).WillReturnRows( pgxmock.NewRows([]string{"id", "name", "canSync"}). AddRow(database.MustToDBUUID(j.ClientID), "client_name", true), ) pool.ExpectQuery("-- name: GetDocumentIDByHash :one").WithArgs(doc.Hash, database.MustToDBUUID(doc.JobID)).WillReturnRows( pgxmock.NewRows([]string{"id"}), ) pool.ExpectBegin() pool.ExpectQuery("name: CreateDocument :one").WithArgs(database.MustToDBUUID(doc.JobID), doc.Hash). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(database.MustToDBUUID(doc.ID)), ) pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(database.MustToDBUUID(doc.ID), bucket, location). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() mockStore.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == location }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ETag: &doc.Hash, }, nil) id, err := svc.Create(ctx, &Create{ JobID: doc.JobID, Location: location, Bucket: bucket, }) assert.NoError(t, err) assert.Equal(t, doc.ID, id) } func TestInformCreate(t *testing.T) { ctx := context.Background() mockSQS := queuemock.NewMockSQSClient(t) cfg := &DocInitConfig{} cfg.QueueClient = mockSQS cfg.DocumentCleanURL = "/i/am/here" svc := Service{ cfg: cfg, } id := uuid.New() j := &job.Job{} err := svc.informCreate(ctx, id, j) assert.NoError(t, err) mockSQS.EXPECT(). SendMessage( mock.Anything, mock.MatchedBy(func(in *sqs.SendMessageInput) bool { return *in.QueueUrl == cfg.DocumentCleanURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id) }), mock.Anything, ). Return(&sqs.SendMessageOutput{}, nil) j.CanSync = true err = svc.informCreate(ctx, id, j) assert.NoError(t, err) } func TestGetCreateParams(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) doc := document.Document{ JobID: uuid.New(), Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectQuery("-- name: GetDocumentIDByHash :one").WithArgs(doc.Hash, database.MustToDBUUID(doc.JobID)).WillReturnRows( pgxmock.NewRows([]string{"id"}), ) mockStore.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == location }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ETag: &doc.Hash, }, nil) params, err := svc.getCreateParams(ctx, &Create{ JobID: doc.JobID, Location: location, Bucket: bucket, }) assert.NoError(t, err) assert.Equal(t, &createDocumentParams{ JobID: database.MustToDBUUID(doc.JobID), Hash: doc.Hash, Bucket: bucket, Location: location, }, params) } func TestGetCreateParamsExisting(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) doc := document.Document{ ID: uuid.New(), JobID: uuid.New(), Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectQuery("-- name: GetDocumentIDByHash :one").WithArgs(doc.Hash, database.MustToDBUUID(doc.JobID)).WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(database.MustToDBUUID(doc.ID)), ) mockStore.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == location }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ETag: &doc.Hash, }, nil) params, err := svc.getCreateParams(ctx, &Create{ JobID: doc.JobID, Location: location, Bucket: bucket, }) assert.NoError(t, err) dbid := database.MustToDBUUID(doc.ID) assert.Equal(t, &createDocumentParams{ ID: &dbid, JobID: database.MustToDBUUID(doc.JobID), Hash: doc.Hash, Bucket: bucket, Location: location, }, params) } func TestGetCreateParamsCurrentDocErr(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) mockStore := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockStore svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) doc := document.Document{ JobID: uuid.New(), Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectQuery("-- name: GetDocumentIDByHash :one").WithArgs(doc.Hash, database.MustToDBUUID(doc.JobID)). WillReturnError(errors.New("db err")) mockStore.EXPECT(). HeadObject( mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool { return *in.Bucket == bucket && *in.Key == location }), mock.Anything, ). Return(&s3.HeadObjectOutput{ ETag: &doc.Hash, }, nil) _, err = svc.getCreateParams(ctx, &Create{ JobID: doc.JobID, Location: location, Bucket: bucket, }) assert.Error(t, err) } func TestSubmitCreate(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) doc := document.Document{ ID: uuid.New(), JobID: uuid.New(), Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectBegin() pool.ExpectQuery("name: CreateDocument :one").WithArgs(database.MustToDBUUID(doc.JobID), doc.Hash). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(database.MustToDBUUID(doc.ID)), ) pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(database.MustToDBUUID(doc.ID), bucket, location). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() id, err := svc.submitCreate(ctx, &createDocumentParams{ JobID: database.MustToDBUUID(doc.JobID), Location: location, Bucket: bucket, Hash: doc.Hash, }) assert.NoError(t, err) assert.Equal(t, doc.ID, id) } func TestSubmitCreateExists(t *testing.T) { ctx := context.Background() pool, err := pgxmock.NewPool() if err != nil { t.Fatalf("failed to open pgxmock database: %v", err) } cfg := &DocInitConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) svc := New(cfg, &Services{ Job: job.New(cfg, &job.Services{ Client: client.New(cfg), }), }) doc := document.Document{ ID: uuid.New(), JobID: uuid.New(), Hash: "example_hash", } bucket := "eample_bucket" location := "/i/am/here" pool.ExpectBegin() pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(database.MustToDBUUID(doc.ID), bucket, location). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() docid := database.MustToDBUUID(doc.ID) id, err := svc.submitCreate(ctx, &createDocumentParams{ ID: &docid, JobID: database.MustToDBUUID(doc.JobID), Location: location, Bucket: bucket, Hash: doc.Hash, }) assert.NoError(t, err) assert.Equal(t, doc.ID, id) }