test: fix existing unit tests

goto · Aug 11, 2024 · 5994448 · 5994448
1 parent 22d6d9e
commit 5994448
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 20 deletions.
diff --git a/core/asset/service_test.go b/core/asset/service_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/goto/compass/core/asset/mocks"
 	"github.com/goto/compass/internal/workermanager"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 )
 
 func TestService_GetAllAssets(t *testing.T) {
@@ -195,7 +196,7 @@ func TestService_UpsertAsset(t *testing.T) {
 			Asset:       sampleAsset,
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
 				ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
-				dr.EXPECT().Upsert(ctx, *sampleAsset).Return(errors.New("unknown error"))
+				dr.EXPECT().Upsert(ctx, mock.AnythingOfType("asset.Asset")).Return(errors.New("unknown error"))
 			},
 			Err:        errors.New("unknown error"),
 			ReturnedID: sampleAsset.ID,
@@ -207,7 +208,7 @@ func TestService_UpsertAsset(t *testing.T) {
 			Downstreams: sampleNodes2,
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
 				ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
-				dr.EXPECT().Upsert(ctx, *sampleAsset).Return(nil)
+				dr.EXPECT().Upsert(ctx, mock.AnythingOfType("asset.Asset")).Return(nil)
 				lr.EXPECT().Upsert(ctx, sampleAsset.URN, sampleNodes1, sampleNodes2).Return(errors.New("unknown error"))
 			},
 			Err:        errors.New("unknown error"),
@@ -220,7 +221,7 @@ func TestService_UpsertAsset(t *testing.T) {
 			Downstreams: sampleNodes2,
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository, lr *mocks.LineageRepository) {
 				ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
-				dr.EXPECT().Upsert(ctx, *sampleAsset).Return(nil)
+				dr.EXPECT().Upsert(ctx, mock.AnythingOfType("asset.Asset")).Return(nil)
 				lr.EXPECT().Upsert(ctx, sampleAsset.URN, sampleNodes1, sampleNodes2).Return(nil)
 			},
 			Err:        nil,
@@ -278,7 +279,7 @@ func TestService_UpsertAssetWithoutLineage(t *testing.T) {
 			Asset:       sampleAsset,
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository) {
 				ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
-				dr.EXPECT().Upsert(ctx, *sampleAsset).Return(errors.New("unknown error"))
+				dr.EXPECT().Upsert(ctx, mock.AnythingOfType("asset.Asset")).Return(errors.New("unknown error"))
 			},
 			Err: errors.New("unknown error"),
 		},
@@ -287,7 +288,7 @@ func TestService_UpsertAssetWithoutLineage(t *testing.T) {
 			Asset:       sampleAsset,
 			Setup: func(ctx context.Context, ar *mocks.AssetRepository, dr *mocks.DiscoveryRepository) {
 				ar.EXPECT().Upsert(ctx, sampleAsset).Return(sampleAsset.ID, nil)
-				dr.EXPECT().Upsert(ctx, *sampleAsset).Return(nil)
+				dr.EXPECT().Upsert(ctx, mock.AnythingOfType("asset.Asset")).Return(nil)
 			},
 			ReturnedID: sampleAsset.ID,
 		},

diff --git a/internal/workermanager/discovery_worker.go b/internal/workermanager/discovery_worker.go
@@ -26,7 +26,7 @@ func (m *Manager) EnqueueIndexAssetJob(ctx context.Context, ast asset.Asset) err
 	}
 
 	err = m.worker.Enqueue(ctx, worker.JobSpec{
-		Type:    jobIndexAsset,
+		Type:    JobIndexAsset,
 		Payload: payload,
 	})
 	if err != nil {
@@ -128,7 +128,7 @@ func (m *Manager) SyncAssets(ctx context.Context, job worker.JobSpec) error {
 
 func (m *Manager) EnqueueDeleteAssetJob(ctx context.Context, urn string) error {
 	err := m.worker.Enqueue(ctx, worker.JobSpec{
-		Type:    jobDeleteAsset,
+		Type:    JobDeleteAsset,
 		Payload: ([]byte)(urn),
 	})
 	if err != nil {
@@ -161,7 +161,7 @@ func (m *Manager) DeleteAsset(ctx context.Context, job worker.JobSpec) error {
 
 func (m *Manager) EnqueueDeleteAssetsByQueryExprJob(ctx context.Context, queryExpr string) error {
 	err := m.worker.Enqueue(ctx, worker.JobSpec{
-		Type:    jobDeleteAssetsByQuery,
+		Type:    JobDeleteAssetsByQuery,
 		Payload: ([]byte)(queryExpr),
 	})
 	if err != nil {
@@ -194,7 +194,7 @@ func (m *Manager) DeleteAssetsByQueryExpr(ctx context.Context, job worker.JobSpe
 
 func (m *Manager) EnqueueSyncAssetJob(ctx context.Context, service string) error {
 	err := m.worker.Enqueue(ctx, worker.JobSpec{
-		Type:    jobSyncAsset,
+		Type:    JobSyncAsset,
 		Payload: ([]byte)(service),
 	})
 	if err != nil {

diff --git a/internal/workermanager/job_types.go b/internal/workermanager/job_types.go
@@ -1,8 +1,8 @@
 package workermanager
 
 const (
-	jobIndexAsset          = "index-asset"
-	jobDeleteAsset         = "delete-asset"
-	jobDeleteAssetsByQuery = "delete-assets-by-query"
-	jobSyncAsset           = "sync-asset"
+	JobIndexAsset          = "index-asset"
+	JobDeleteAsset         = "delete-asset"
+	JobDeleteAssetsByQuery = "delete-assets-by-query"
+	JobSyncAsset           = "sync-asset"
 )
diff --git a/internal/workermanager/worker_manager.go b/internal/workermanager/worker_manager.go
@@ -124,10 +124,10 @@ func (m *Manager) init() error {
 	m.initDone.Store(true)
 
 	jobHandlers := map[string]worker.JobHandler{
-		jobIndexAsset:          m.indexAssetHandler(),
-		jobDeleteAsset:         m.deleteAssetHandler(),
-		jobDeleteAssetsByQuery: m.deleteAssetsByQueryHandler(),
-		jobSyncAsset:           m.syncAssetHandler(),
+		JobIndexAsset:          m.indexAssetHandler(),
+		JobDeleteAsset:         m.deleteAssetHandler(),
+		JobDeleteAssetsByQuery: m.deleteAssetsByQueryHandler(),
+		JobSyncAsset:           m.syncAssetHandler(),
 	}
 	for typ, h := range jobHandlers {
 		if err := m.worker.Register(typ, h); err != nil {

diff --git a/internal/workermanager/worker_manager_test.go b/internal/workermanager/worker_manager_test.go
@@ -77,13 +77,16 @@ func TestManager_Run(t *testing.T) {
 			wrkr := mocks.NewWorker(t)
 			mgr := workermanager.NewWithWorker(wrkr, workermanager.Deps{})
 			wrkr.EXPECT().
-				Register("index-asset", mock.AnythingOfType("worker.JobHandler")).
+				Register(workermanager.JobIndexAsset, mock.AnythingOfType("worker.JobHandler")).
 				Return(nil)
 			wrkr.EXPECT().
-				Register("delete-asset", mock.AnythingOfType("worker.JobHandler")).
+				Register(workermanager.JobDeleteAsset, mock.AnythingOfType("worker.JobHandler")).
 				Return(nil)
 			wrkr.EXPECT().
-				Register("sync-asset", mock.AnythingOfType("worker.JobHandler")).
+				Register(workermanager.JobDeleteAssetsByQuery, mock.AnythingOfType("worker.JobHandler")).
+				Return(nil)
+			wrkr.EXPECT().
+				Register(workermanager.JobSyncAsset, mock.AnythingOfType("worker.JobHandler")).
 				Return(nil)
 			wrkr.EXPECT().
 				Run(ctx).