use mockery to generate mocks (#6244)

Signed-off-by: Alex Wu <[email protected]>
flyteorg · Feb 14, 2025 · 0634811 · 0634811
1 parent 503db30
commit 0634811
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 52 deletions.
diff --git a/flyteadmin/pkg/async/cloudevent/factory_test.go b/flyteadmin/pkg/async/cloudevent/factory_test.go
@@ -42,7 +42,7 @@ func TestGetCloudEventPublisher(t *testing.T) {
 
 	db := mocks.NewMockRepository()
 	mockStore := getMockStore()
-	url := dataMocks.NewMockRemoteURL()
+	url := &dataMocks.RemoteURLInterface{}
 
 	t.Run("local publisher", func(t *testing.T) {
 		cfg.Type = common.Local
@@ -64,7 +64,7 @@ func TestInvalidAwsConfig(t *testing.T) {
 	}
 	db := mocks.NewMockRepository()
 	mockStore := getMockStore()
-	url := dataMocks.NewMockRemoteURL()
+	url := &dataMocks.RemoteURLInterface{}
 
 	NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
 	t.Errorf("did not panic")
@@ -79,7 +79,7 @@ func TestInvalidGcpConfig(t *testing.T) {
 	}
 	db := mocks.NewMockRepository()
 	mockStore := getMockStore()
-	url := dataMocks.NewMockRemoteURL()
+	url := &dataMocks.RemoteURLInterface{}
 
 	NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
 	t.Errorf("did not panic")
@@ -95,7 +95,7 @@ func TestInvalidKafkaConfig(t *testing.T) {
 	}
 	db := mocks.NewMockRepository()
 	mockStore := getMockStore()
-	url := dataMocks.NewMockRemoteURL()
+	url := &dataMocks.RemoteURLInterface{}
 
 	NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
 	cfg.KafkaConfig = runtimeInterfaces.KafkaConfig{Version: "2.1.0"}

diff --git a/flyteadmin/pkg/data/interfaces/remote.go b/flyteadmin/pkg/data/interfaces/remote.go
@@ -6,6 +6,8 @@ import (
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
 )
 
+//go:generate mockery-v2 --name=RemoteURLInterface --output=../mocks --case=underscore --with-expecter
+
 // Defines an interface for fetching pre-signed URLs.
 type RemoteURLInterface interface {
 	// TODO: Refactor for URI to be of type DataReference. We should package a FromString-like function in flytestdlib

diff --git a/flyteadmin/pkg/data/mocks/remote.go b/flyteadmin/pkg/data/mocks/remote.go
diff --git a/flyteadmin/pkg/data/mocks/remote_url_interface.go b/flyteadmin/pkg/data/mocks/remote_url_interface.go
diff --git a/flyteadmin/pkg/manager/impl/execution_manager_test.go b/flyteadmin/pkg/manager/impl/execution_manager_test.go
@@ -79,7 +79,7 @@ var executionIdentifier = core.WorkflowExecutionIdentifier{
 	Name:    "name",
 }
 var mockPublisher notificationMocks.Publisher
-var mockExecutionRemoteURL = dataMocks.NewMockRemoteURL()
+var mockExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
 var requestedAt = time.Now()
 var testCluster = "C1"
 var outputURI = "output uri"
@@ -3624,8 +3624,8 @@ func TestGetExecutionData(t *testing.T) {
 			InputsURI:    shared.Inputs,
 		}, nil
 	}
-	mockExecutionRemoteURL := dataMocks.NewMockRemoteURL()
-	mockExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(
+	mockExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
+	mockExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(
 		ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		if uri == outputURI {
 			return &admin.UrlBlob{
@@ -3640,7 +3640,7 @@ func TestGetExecutionData(t *testing.T) {
 		}
 
 		return &admin.UrlBlob{}, errors.New("unexpected input")
-	}
+	})
 	mockStorage := commonMocks.GetMockStorageClient()
 	fullInputs := &core.LiteralMap{
 		Literals: map[string]*core.Literal{
@@ -3842,8 +3842,8 @@ func TestGetExecutionData_LegacyModel(t *testing.T) {
 			StartedAt:    &startedAt,
 		}, nil
 	}
-	mockExecutionRemoteURL := dataMocks.NewMockRemoteURL()
-	mockExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(
+	mockExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
+	mockExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(
 		ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		if uri == outputURI {
 			return &admin.UrlBlob{
@@ -3858,7 +3858,7 @@ func TestGetExecutionData_LegacyModel(t *testing.T) {
 		}
 
 		return &admin.UrlBlob{}, errors.New("unexpected input")
-	}
+	})
 
 	repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetGetCallback(executionGetFunc)
 	storageClient := getMockStorageForExecTest(context.Background())

diff --git a/flyteadmin/pkg/manager/impl/node_execution_manager_test.go b/flyteadmin/pkg/manager/impl/node_execution_manager_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/ptypes"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -97,7 +98,7 @@ var workflowExecutionIdentifier = core.WorkflowExecutionIdentifier{
 	Name:    "name",
 }
 
-var mockNodeExecutionRemoteURL = dataMocks.NewMockRemoteURL()
+var mockNodeExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
 
 func addGetExecutionCallback(t *testing.T, repository interfaces.Repository) {
 	repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetGetCallback(
@@ -1377,8 +1378,8 @@ func TestGetNodeExecutionData(t *testing.T) {
 			}, nil
 		})
 
-	mockNodeExecutionRemoteURL := dataMocks.NewMockRemoteURL()
-	mockNodeExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
+	mockNodeExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
+	mockNodeExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		if uri == "input uri" {
 			return &admin.UrlBlob{
 				Url:   "inputs",
@@ -1392,7 +1393,7 @@ func TestGetNodeExecutionData(t *testing.T) {
 		}
 
 		return &admin.UrlBlob{}, errors.New("unexpected input")
-	}
+	})
 	mockStorage := commonMocks.GetMockStorageClient()
 	fullInputs := &core.LiteralMap{
 		Literals: map[string]*core.Literal{

diff --git a/flyteadmin/pkg/manager/impl/task_execution_manager_test.go b/flyteadmin/pkg/manager/impl/task_execution_manager_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	"github.com/golang/protobuf/ptypes"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc/codes"
 
 	"github.com/flyteorg/flyte/flyteadmin/pkg/common"
@@ -63,7 +64,7 @@ var taskEventRequest = &admin.TaskExecutionEventRequest{
 	},
 }
 
-var mockTaskExecutionRemoteURL = dataMocks.NewMockRemoteURL()
+var mockTaskExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
 
 var retryAttemptValue = uint32(1)
 
@@ -1070,8 +1071,8 @@ func TestGetTaskExecutionData(t *testing.T) {
 				Closure:   closureBytes,
 			}, nil
 		})
-	mockTaskExecutionRemoteURL = dataMocks.NewMockRemoteURL()
-	mockTaskExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
+	mockTaskExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
+	mockTaskExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		if uri == "input-uri.pb" {
 			return &admin.UrlBlob{
 				Url:   "inputs",
@@ -1085,7 +1086,7 @@ func TestGetTaskExecutionData(t *testing.T) {
 		}
 
 		return &admin.UrlBlob{}, errors.New("unexpected input")
-	}
+	})
 	mockStorage := commonMocks.GetMockStorageClient()
 	fullInputs := &core.LiteralMap{
 		Literals: map[string]*core.Literal{

diff --git a/flyteadmin/pkg/manager/impl/util/data_test.go b/flyteadmin/pkg/manager/impl/util/data_test.go
@@ -6,6 +6,7 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/flyteorg/flyte/flyteadmin/pkg/common"
 	commonMocks "github.com/flyteorg/flyte/flyteadmin/pkg/common/mocks"
@@ -118,11 +119,11 @@ func TestGetInputs(t *testing.T) {
 		Bytes: 1000,
 	}
 
-	mockRemoteURL := urlMocks.NewMockRemoteURL()
-	mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
+	mockRemoteURL := &urlMocks.RemoteURLInterface{}
+	mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		assert.Equal(t, inputsURI, uri)
 		return expectedURLBlob, nil
-	}
+	})
 	remoteDataConfig := interfaces.RemoteDataConfig{
 		MaxSizeInBytes: 2000,
 	}
@@ -162,11 +163,11 @@ func TestGetOutputs(t *testing.T) {
 		Bytes: 1000,
 	}
 
-	mockRemoteURL := urlMocks.NewMockRemoteURL()
-	mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
+	mockRemoteURL := &urlMocks.RemoteURLInterface{}
+	mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
 		assert.Equal(t, testOutputsURI, uri)
 		return expectedURLBlob, nil
-	}
+	})
 
 	remoteDataConfig := interfaces.RemoteDataConfig{
 		MaxSizeInBytes: 2000,
@@ -205,11 +206,11 @@ func TestGetOutputs(t *testing.T) {
 		assert.Empty(t, outputURLBlob)
 	})
 	t.Run("inline outputs", func(t *testing.T) {
-		mockRemoteURL := urlMocks.NewMockRemoteURL()
-		mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
+		mockRemoteURL := &urlMocks.RemoteURLInterface{}
+		mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
 			t.Fatal("Should not fetch a remote URL for outputs stored inline for an execution model")
 			return &admin.UrlBlob{}, nil
-		}
+		})
 		remoteDataConfig := interfaces.RemoteDataConfig{}
 		remoteDataConfig.MaxSizeInBytes = 2000