Switch workflowengine compiler mock to mockery generated

flyteadmin/pkg/manager/impl/task_manager_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc/codes"
 
 	"github.com/flyteorg/flyte/flyteadmin/pkg/common"
@@ -52,14 +53,12 @@ func init() {
 }
 
 func getMockTaskCompiler() workflowengine.Compiler {
-	mockCompiler := workflowMocks.NewMockCompiler()
-	mockCompiler.(*workflowMocks.MockCompiler).AddCompileTaskCallback(
-		func(task *core.TaskTemplate) (*core.CompiledTask, error) {
-			return &core.CompiledTask{
-				Template: task,
-			}, nil
-		})
-	return mockCompiler
+	compiler := &workflowMocks.Compiler{}
+	compiler.On("CompileTask", mock.AnythingOfType("*core.TaskTemplate")).Return(func(taskTemplate *core.TaskTemplate) (*core.CompiledTask, error) {
+		return &core.CompiledTask{Template: taskTemplate}, nil
+	})
+
+	return compiler
 }
 
 func getMockTaskRepository() interfaces.Repository {
@@ -172,12 +171,10 @@ func TestCreateTask_ValidationError(t *testing.T) {
 }
 
 func TestCreateTask_CompilerError(t *testing.T) {
-	mockCompiler := workflowMocks.NewMockCompiler()
+	mockCompiler := &workflowMocks.Compiler{}
 	expectedErr := errors.New("expected error")
-	mockCompiler.(*workflowMocks.MockCompiler).AddCompileTaskCallback(
-		func(task *core.TaskTemplate) (*core.CompiledTask, error) {
-			return nil, expectedErr
-		})
+	mockCompiler.On("CompileTask", mock.MatchedBy(func(*core.TaskTemplate) bool { return true })).Return(nil, expectedErr)
+
 	mockRepository := getMockTaskRepository()
 	taskManager := NewTaskManager(mockRepository, getMockConfigForTaskTest(), mockCompiler,
 		mockScope.NewTestScope())

flyteadmin/pkg/manager/impl/workflow_manager_test.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -89,21 +91,19 @@ func getMockRepository(workflowOnGet bool) interfaces.Repository {
 }
 
 func getMockWorkflowCompiler() workflowengineInterfaces.Compiler {
-	mockCompiler := workflowengineMocks.NewMockCompiler()
-	mockCompiler.(*workflowengineMocks.MockCompiler).AddGetRequirementCallback(
-		func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) (
-			reqs compiler.WorkflowExecutionRequirements, err error) {
-			return compiler.WorkflowExecutionRequirements{}, nil
-		})
-	mockCompiler.(*workflowengineMocks.MockCompiler).AddCompileWorkflowCallback(func(
-		primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask,
-		launchPlans []engine.InterfaceProvider) (*core.CompiledWorkflowClosure, error) {
+	mockCompiler := &workflowengineMocks.Compiler{}
+	mockCompiler.On("GetRequirements", mock.AnythingOfType("*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.WorkflowTemplate")).Return(func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) (compiler.WorkflowExecutionRequirements, error) {
+		return compiler.WorkflowExecutionRequirements{}, nil
+	})
+
+	mockCompiler.On("CompileWorkflow", mock.AnythingOfType("*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.CompiledTask"), mock.AnythingOfType("[]common.InterfaceProvider")).Return(func(primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, launchPlans []engine.InterfaceProvider) (*core.CompiledWorkflowClosure, error) {
 		return &core.CompiledWorkflowClosure{
 			Primary: &core.CompiledWorkflow{
 				Template: primaryWf,
 			},
 		}, nil
 	})
+
 	return mockCompiler
 }
 
@@ -213,12 +213,10 @@ func TestCreateWorkflow_ExistingWorkflow_Different(t *testing.T) {
 
 func TestCreateWorkflow_CompilerGetRequirementsError(t *testing.T) {
 	expectedErr := errors.New("expected error")
-	mockCompiler := getMockWorkflowCompiler()
-	mockCompiler.(*workflowengineMocks.MockCompiler).AddGetRequirementCallback(
-		func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) (
-			reqs compiler.WorkflowExecutionRequirements, err error) {
-			return compiler.WorkflowExecutionRequirements{}, expectedErr
-		})
+	mockCompiler := &workflowengineMocks.Compiler{}
+	mockCompiler.On("GetRequirements", mock.AnythingOfType("*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.WorkflowTemplate")).Return(func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) (compiler.WorkflowExecutionRequirements, error) {
+		return compiler.WorkflowExecutionRequirements{}, expectedErr
+	})
 
 	workflowManager := NewWorkflowManager(
 		getMockRepository(!returnWorkflowOnGet),
@@ -233,16 +231,16 @@ func TestCreateWorkflow_CompilerGetRequirementsError(t *testing.T) {
 
 func TestCreateWorkflow_CompileWorkflowError(t *testing.T) {
 	expectedErr := errors.New("expected error")
-	mockCompiler := getMockWorkflowCompiler()
-	mockCompiler.(*workflowengineMocks.MockCompiler).AddCompileWorkflowCallback(func(
-		primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask,
-		launchPlans []engine.InterfaceProvider) (*core.CompiledWorkflowClosure, error) {
+	mockCompiler := workflowengineMocks.Compiler{}
+	mockCompiler.On("GetRequirements", mock.AnythingOfType("*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.WorkflowTemplate")).Return(compiler.WorkflowExecutionRequirements{}, nil)
+
+	mockCompiler.On("CompileWorkflow", mock.AnythingOfType("*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.WorkflowTemplate"), mock.AnythingOfType("[]*core.CompiledTask"), mock.AnythingOfType("[]common.InterfaceProvider")).Return(func(primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, launchPlans []engine.InterfaceProvider) (*core.CompiledWorkflowClosure, error) {
 		return &core.CompiledWorkflowClosure{}, expectedErr
 	})
 
 	workflowManager := NewWorkflowManager(
 		getMockRepository(!returnWorkflowOnGet),
-		getMockWorkflowConfigProvider(), mockCompiler, getMockStorage(), storagePrefix, mockScope.NewTestScope())
+		getMockWorkflowConfigProvider(), &mockCompiler, getMockStorage(), storagePrefix, mockScope.NewTestScope())
 	request := testutils.GetWorkflowRequest()
 	response, err := workflowManager.CreateWorkflow(context.Background(), request)
 	assert.Nil(t, response)

flyteadmin/pkg/workflowengine/interfaces/builder.go

@@ -5,7 +5,7 @@ import (
 	"github.com/flyteorg/flyte/flytepropeller/pkg/apis/flyteworkflow/v1alpha1"
 )
 
-//go:generate mockery -name FlyteWorkflowBuilder -output=../mocks -case=underscore
+//go:generate mockery-v2 --name=FlyteWorkflowBuilder --output=../mocks --case=underscore
 
 // FlyteWorkflowBuilder produces a v1alpha1.FlyteWorkflow definition from a compiled workflow closure and execution inputs
 type FlyteWorkflowBuilder interface {

flyteadmin/pkg/workflowengine/interfaces/compiler.go

@@ -6,6 +6,8 @@ import (
 	"github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common"
 )
 
+//go:generate mockery-v2 --name=Compiler --output=../mocks --case=underscore
+
 // Workflow compiler interface.
 type Compiler interface {
 	CompileTask(task *core.TaskTemplate) (*core.CompiledTask, error)

flyteadmin/pkg/workflowengine/interfaces/executor.go

@@ -10,7 +10,7 @@ import (
 	"github.com/flyteorg/flyte/flytestdlib/storage"
 )
 
-//go:generate mockery -name=WorkflowExecutor -output=../mocks/ -case=underscore
+//go:generate mockery-v2 --name=WorkflowExecutor --output=../mocks --case=underscore
 
 type TaskResources struct {
 	Defaults runtime.TaskResourceSet