From f12185a95b11a9d4e1a0aac3107f0ee8d9ce61ae Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 19 Feb 2025 11:40:37 +0800 Subject: [PATCH 1/5] remove old hand roll mock and add mockery generated compiler mock Signed-off-by: Arthur --- .../pkg/workflowengine/mocks/Compiler.go | 119 ++++++++++++++++++ .../pkg/workflowengine/mocks/mock_compiler.go | 62 --------- 2 files changed, 119 insertions(+), 62 deletions(-) create mode 100644 flyteadmin/pkg/workflowengine/mocks/Compiler.go delete mode 100644 flyteadmin/pkg/workflowengine/mocks/mock_compiler.go diff --git a/flyteadmin/pkg/workflowengine/mocks/Compiler.go b/flyteadmin/pkg/workflowengine/mocks/Compiler.go new file mode 100644 index 0000000000..bb9a87ae9a --- /dev/null +++ b/flyteadmin/pkg/workflowengine/mocks/Compiler.go @@ -0,0 +1,119 @@ +// Code generated by mockery v2.40.3. DO NOT EDIT. + +package mocks + +import ( + compiler "github.com/flyteorg/flyte/flytepropeller/pkg/compiler" + common "github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common" + + core "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" + + mock "github.com/stretchr/testify/mock" +) + +// Compiler is an autogenerated mock type for the Compiler type +type Compiler struct { + mock.Mock +} + +// CompileTask provides a mock function with given fields: task +func (_m *Compiler) CompileTask(task *core.TaskTemplate) (*core.CompiledTask, error) { + ret := _m.Called(task) + + if len(ret) == 0 { + panic("no return value specified for CompileTask") + } + + var r0 *core.CompiledTask + var r1 error + if rf, ok := ret.Get(0).(func(*core.TaskTemplate) (*core.CompiledTask, error)); ok { + return rf(task) + } + if rf, ok := ret.Get(0).(func(*core.TaskTemplate) *core.CompiledTask); ok { + r0 = rf(task) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*core.CompiledTask) + } + } + + if rf, ok := ret.Get(1).(func(*core.TaskTemplate) error); ok { + r1 = rf(task) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// CompileWorkflow provides a mock function with given fields: primaryWf, subworkflows, tasks, launchPlans +func (_m *Compiler) CompileWorkflow(primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, launchPlans []common.InterfaceProvider) (*core.CompiledWorkflowClosure, error) { + ret := _m.Called(primaryWf, subworkflows, tasks, launchPlans) + + if len(ret) == 0 { + panic("no return value specified for CompileWorkflow") + } + + var r0 *core.CompiledWorkflowClosure + var r1 error + if rf, ok := ret.Get(0).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate, []*core.CompiledTask, []common.InterfaceProvider) (*core.CompiledWorkflowClosure, error)); ok { + return rf(primaryWf, subworkflows, tasks, launchPlans) + } + if rf, ok := ret.Get(0).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate, []*core.CompiledTask, []common.InterfaceProvider) *core.CompiledWorkflowClosure); ok { + r0 = rf(primaryWf, subworkflows, tasks, launchPlans) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*core.CompiledWorkflowClosure) + } + } + + if rf, ok := ret.Get(1).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate, []*core.CompiledTask, []common.InterfaceProvider) error); ok { + r1 = rf(primaryWf, subworkflows, tasks, launchPlans) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// GetRequirements provides a mock function with given fields: fg, subWfs +func (_m *Compiler) GetRequirements(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) (compiler.WorkflowExecutionRequirements, error) { + ret := _m.Called(fg, subWfs) + + if len(ret) == 0 { + panic("no return value specified for GetRequirements") + } + + var r0 compiler.WorkflowExecutionRequirements + var r1 error + if rf, ok := ret.Get(0).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate) (compiler.WorkflowExecutionRequirements, error)); ok { + return rf(fg, subWfs) + } + if rf, ok := ret.Get(0).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate) compiler.WorkflowExecutionRequirements); ok { + r0 = rf(fg, subWfs) + } else { + r0 = ret.Get(0).(compiler.WorkflowExecutionRequirements) + } + + if rf, ok := ret.Get(1).(func(*core.WorkflowTemplate, []*core.WorkflowTemplate) error); ok { + r1 = rf(fg, subWfs) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// NewCompiler creates a new instance of Compiler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewCompiler(t interface { + mock.TestingT + Cleanup(func()) +}) *Compiler { + mock := &Compiler{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/flyteadmin/pkg/workflowengine/mocks/mock_compiler.go b/flyteadmin/pkg/workflowengine/mocks/mock_compiler.go deleted file mode 100644 index eaf27ea433..0000000000 --- a/flyteadmin/pkg/workflowengine/mocks/mock_compiler.go +++ /dev/null @@ -1,62 +0,0 @@ -package mocks - -import ( - "github.com/flyteorg/flyte/flyteadmin/pkg/workflowengine/interfaces" - "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" - "github.com/flyteorg/flyte/flytepropeller/pkg/compiler" - "github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common" -) - -type MockCompiler struct { - compileTaskCallback func(task *core.TaskTemplate) (*core.CompiledTask, error) - getRequirementsCallback func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) ( - reqs compiler.WorkflowExecutionRequirements, err error) - compileWorkflowCallback func( - primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, - launchPlans []common.InterfaceProvider) (*core.CompiledWorkflowClosure, error) -} - -func (c *MockCompiler) CompileTask(task *core.TaskTemplate) (*core.CompiledTask, error) { - if c.compileTaskCallback != nil { - return c.compileTaskCallback(task) - } - return nil, nil -} - -func (c *MockCompiler) AddCompileTaskCallback( - callback func(task *core.TaskTemplate) (*core.CompiledTask, error)) { - c.compileTaskCallback = callback -} - -func (c *MockCompiler) GetRequirements(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) ( - reqs compiler.WorkflowExecutionRequirements, err error) { - if c.getRequirementsCallback != nil { - return c.getRequirementsCallback(fg, subWfs) - } - return compiler.WorkflowExecutionRequirements{}, nil -} - -func (c *MockCompiler) AddGetRequirementCallback( - callback func(fg *core.WorkflowTemplate, subWfs []*core.WorkflowTemplate) ( - reqs compiler.WorkflowExecutionRequirements, err error)) { - c.getRequirementsCallback = callback -} - -func (c *MockCompiler) CompileWorkflow( - primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, - launchPlans []common.InterfaceProvider) (*core.CompiledWorkflowClosure, error) { - if c.compileWorkflowCallback != nil { - return c.compileWorkflowCallback(primaryWf, subworkflows, tasks, launchPlans) - } - return nil, nil -} - -func (c *MockCompiler) AddCompileWorkflowCallback(callback func( - primaryWf *core.WorkflowTemplate, subworkflows []*core.WorkflowTemplate, tasks []*core.CompiledTask, - launchPlans []common.InterfaceProvider) (*core.CompiledWorkflowClosure, error)) { - c.compileWorkflowCallback = callback -} - -func NewMockCompiler() interfaces.Compiler { - return &MockCompiler{} -} From 0ea8e33047a48474d9081457b5fe4f48aa8c7530 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 19 Feb 2025 15:54:30 +0800 Subject: [PATCH 2/5] use new mockery generated workflowengine compiler mock in task_manager_test the AddCompileTaskCallback is only used in old hand roll mock, nowhere else is using it, fine to remove it Signed-off-by: Arthur --- .../pkg/manager/impl/task_manager_test.go | 16 ++-------- .../pkg/manager/impl/workflow_manager_test.go | 29 ++----------------- 2 files changed, 5 insertions(+), 40 deletions(-) diff --git a/flyteadmin/pkg/manager/impl/task_manager_test.go b/flyteadmin/pkg/manager/impl/task_manager_test.go index 50fc2c3234..1431478a80 100644 --- a/flyteadmin/pkg/manager/impl/task_manager_test.go +++ b/flyteadmin/pkg/manager/impl/task_manager_test.go @@ -52,14 +52,7 @@ 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 + return &workflowMocks.Compiler{} } func getMockTaskRepository() interfaces.Repository { @@ -172,12 +165,9 @@ 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 - }) + mockRepository := getMockTaskRepository() taskManager := NewTaskManager(mockRepository, getMockConfigForTaskTest(), mockCompiler, mockScope.NewTestScope()) diff --git a/flyteadmin/pkg/manager/impl/workflow_manager_test.go b/flyteadmin/pkg/manager/impl/workflow_manager_test.go index 280624a956..d293e1b374 100644 --- a/flyteadmin/pkg/manager/impl/workflow_manager_test.go +++ b/flyteadmin/pkg/manager/impl/workflow_manager_test.go @@ -25,7 +25,6 @@ import ( workflowengineMocks "github.com/flyteorg/flyte/flyteadmin/pkg/workflowengine/mocks" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" - "github.com/flyteorg/flyte/flytepropeller/pkg/compiler" engine "github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common" mockScope "github.com/flyteorg/flyte/flytestdlib/promutils" "github.com/flyteorg/flyte/flytestdlib/storage" @@ -89,22 +88,8 @@ 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) { - return &core.CompiledWorkflowClosure{ - Primary: &core.CompiledWorkflow{ - Template: primaryWf, - }, - }, nil - }) - return mockCompiler + return &workflowengineMocks.Compiler{} + } func getMockStorage() *storage.DataStore { @@ -214,11 +199,6 @@ 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 - }) workflowManager := NewWorkflowManager( getMockRepository(!returnWorkflowOnGet), @@ -234,11 +214,6 @@ 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) { - return &core.CompiledWorkflowClosure{}, expectedErr - }) workflowManager := NewWorkflowManager( getMockRepository(!returnWorkflowOnGet), From 6fbb5a6913b54968854c927db50c2406c315f9af Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 20 Feb 2025 14:35:40 +0800 Subject: [PATCH 3/5] update unit test to use mockery generated mock compiler + add go generate into compiler Signed-off-by: Arthur --- .../pkg/manager/impl/task_manager_test.go | 12 +++++-- .../pkg/manager/impl/workflow_manager_test.go | 33 ++++++++++++++++--- .../pkg/workflowengine/interfaces/compiler.go | 2 ++ 3 files changed, 41 insertions(+), 6 deletions(-) diff --git a/flyteadmin/pkg/manager/impl/task_manager_test.go b/flyteadmin/pkg/manager/impl/task_manager_test.go index 1431478a80..8152b2706a 100644 --- a/flyteadmin/pkg/manager/impl/task_manager_test.go +++ b/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,7 +53,12 @@ func init() { } func getMockTaskCompiler() workflowengine.Compiler { - return &workflowMocks.Compiler{} + 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 { @@ -167,7 +173,9 @@ func TestCreateTask_ValidationError(t *testing.T) { func TestCreateTask_CompilerError(t *testing.T) { mockCompiler := &workflowMocks.Compiler{} expectedErr := errors.New("expected error") - + mockCompiler.On("CompileTask", mock.AnythingOfType("*core.TaskTemplate")).Return(func(taskTemplate *core.TaskTemplate) (*core.CompiledTask, error) { + return nil, expectedErr + }) mockRepository := getMockTaskRepository() taskManager := NewTaskManager(mockRepository, getMockConfigForTaskTest(), mockCompiler, mockScope.NewTestScope()) diff --git a/flyteadmin/pkg/manager/impl/workflow_manager_test.go b/flyteadmin/pkg/manager/impl/workflow_manager_test.go index d293e1b374..994ad7db9a 100644 --- a/flyteadmin/pkg/manager/impl/workflow_manager_test.go +++ b/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" @@ -25,6 +27,7 @@ import ( workflowengineMocks "github.com/flyteorg/flyte/flyteadmin/pkg/workflowengine/mocks" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin" "github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core" + "github.com/flyteorg/flyte/flytepropeller/pkg/compiler" engine "github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common" mockScope "github.com/flyteorg/flyte/flytestdlib/promutils" "github.com/flyteorg/flyte/flytestdlib/storage" @@ -88,8 +91,20 @@ func getMockRepository(workflowOnGet bool) interfaces.Repository { } func getMockWorkflowCompiler() workflowengineInterfaces.Compiler { - return &workflowengineMocks.Compiler{} + 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 } func getMockStorage() *storage.DataStore { @@ -198,7 +213,10 @@ func TestCreateWorkflow_ExistingWorkflow_Different(t *testing.T) { func TestCreateWorkflow_CompilerGetRequirementsError(t *testing.T) { expectedErr := errors.New("expected error") - mockCompiler := getMockWorkflowCompiler() + 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), @@ -213,11 +231,18 @@ func TestCreateWorkflow_CompilerGetRequirementsError(t *testing.T) { func TestCreateWorkflow_CompileWorkflowError(t *testing.T) { expectedErr := errors.New("expected error") - mockCompiler := getMockWorkflowCompiler() + 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{}, 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) diff --git a/flyteadmin/pkg/workflowengine/interfaces/compiler.go b/flyteadmin/pkg/workflowengine/interfaces/compiler.go index 5acc4efa4b..ca263267ab 100644 --- a/flyteadmin/pkg/workflowengine/interfaces/compiler.go +++ b/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 --with-expecter + // Workflow compiler interface. type Compiler interface { CompileTask(task *core.TaskTemplate) (*core.CompiledTask, error) From 5db164d1fda818271d4e25a03071bcd89df4b957 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 20 Feb 2025 17:16:06 +0800 Subject: [PATCH 4/5] update go generate command in workflowengine Signed-off-by: Arthur --- flyteadmin/pkg/workflowengine/interfaces/builder.go | 2 +- flyteadmin/pkg/workflowengine/interfaces/compiler.go | 2 +- flyteadmin/pkg/workflowengine/interfaces/executor.go | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flyteadmin/pkg/workflowengine/interfaces/builder.go b/flyteadmin/pkg/workflowengine/interfaces/builder.go index 7fee0ba131..ba0277e929 100644 --- a/flyteadmin/pkg/workflowengine/interfaces/builder.go +++ b/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 { diff --git a/flyteadmin/pkg/workflowengine/interfaces/compiler.go b/flyteadmin/pkg/workflowengine/interfaces/compiler.go index ca263267ab..e3307ddfeb 100644 --- a/flyteadmin/pkg/workflowengine/interfaces/compiler.go +++ b/flyteadmin/pkg/workflowengine/interfaces/compiler.go @@ -6,7 +6,7 @@ import ( "github.com/flyteorg/flyte/flytepropeller/pkg/compiler/common" ) -//go:generate mockery-v2 --name=Compiler --output=../mocks --case=underscore --with-expecter +//go:generate mockery-v2 --name=Compiler --output=../mocks --case=underscore // Workflow compiler interface. type Compiler interface { diff --git a/flyteadmin/pkg/workflowengine/interfaces/executor.go b/flyteadmin/pkg/workflowengine/interfaces/executor.go index 27de882a5f..463e296948 100644 --- a/flyteadmin/pkg/workflowengine/interfaces/executor.go +++ b/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 From 23fe5772a05efc24fb5b4e8700fe512763bfe07b Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 20 Feb 2025 17:47:02 +0800 Subject: [PATCH 5/5] improve mock function usage Signed-off-by: Arthur --- flyteadmin/pkg/manager/impl/task_manager_test.go | 5 ++--- flyteadmin/pkg/manager/impl/workflow_manager_test.go | 4 +--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/flyteadmin/pkg/manager/impl/task_manager_test.go b/flyteadmin/pkg/manager/impl/task_manager_test.go index 8152b2706a..bbb9aeec7d 100644 --- a/flyteadmin/pkg/manager/impl/task_manager_test.go +++ b/flyteadmin/pkg/manager/impl/task_manager_test.go @@ -173,9 +173,8 @@ func TestCreateTask_ValidationError(t *testing.T) { func TestCreateTask_CompilerError(t *testing.T) { mockCompiler := &workflowMocks.Compiler{} expectedErr := errors.New("expected error") - mockCompiler.On("CompileTask", mock.AnythingOfType("*core.TaskTemplate")).Return(func(taskTemplate *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()) diff --git a/flyteadmin/pkg/manager/impl/workflow_manager_test.go b/flyteadmin/pkg/manager/impl/workflow_manager_test.go index 994ad7db9a..87c98c5cdf 100644 --- a/flyteadmin/pkg/manager/impl/workflow_manager_test.go +++ b/flyteadmin/pkg/manager/impl/workflow_manager_test.go @@ -232,9 +232,7 @@ func TestCreateWorkflow_CompilerGetRequirementsError(t *testing.T) { func TestCreateWorkflow_CompileWorkflowError(t *testing.T) { expectedErr := errors.New("expected 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("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