use mockery generate mock (#6243)

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

Hide file tree

Showing 3 changed files with 88 additions and 30 deletions.
diff --git a/flyteadmin/pkg/clusterresource/controller_test.go b/flyteadmin/pkg/clusterresource/controller_test.go
@@ -121,7 +121,7 @@ func TestPopulateDefaultTemplateValues(t *testing.T) {
 
 func TestGetCustomTemplateValues(t *testing.T) {
 	adminDataProvider := mocks.FlyteAdminDataProvider{}
-	adminDataProvider.OnGetClusterResourceAttributesMatch(mock.Anything, proj, domain).Return(&admin.ClusterResourceAttributes{
+	adminDataProvider.EXPECT().GetClusterResourceAttributes(mock.Anything, mock.Anything, mock.Anything).Return(&admin.ClusterResourceAttributes{
 		Attributes: map[string]string{
 			"var1": "val1",
 			"var2": "val2",
@@ -148,7 +148,7 @@ func TestGetCustomTemplateValues(t *testing.T) {
 
 func TestGetCustomTemplateValues_NothingToOverride(t *testing.T) {
 	adminDataProvider := mocks.FlyteAdminDataProvider{}
-	adminDataProvider.OnGetClusterResourceAttributesMatch(mock.Anything, proj, domain).Return(
+	adminDataProvider.EXPECT().GetClusterResourceAttributes(mock.Anything, mock.Anything, mock.Anything).Return(
 		nil, errors.NewFlyteAdminError(codes.NotFound, "foo"))
 	testController := controller{
 		adminDataProvider: &adminDataProvider,
@@ -332,7 +332,7 @@ imagePullSecrets:
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			adminDataProvider := mocks.FlyteAdminDataProvider{}
-			adminDataProvider.OnGetClusterResourceAttributesMatch(mock.Anything, mock.Anything, mock.Anything).Return(&admin.ClusterResourceAttributes{}, nil)
+			adminDataProvider.EXPECT().GetClusterResourceAttributes(mock.Anything, mock.Anything, mock.Anything).Return(&admin.ClusterResourceAttributes{}, nil)
 			mockPromScope := mockScope.NewTestScope()
 
 			c := NewClusterResourceController(&adminDataProvider, &execClusterMocks.ListTargetsInterface{}, mockPromScope)

diff --git a/flyteadmin/pkg/clusterresource/interfaces/admin.go b/flyteadmin/pkg/clusterresource/interfaces/admin.go
@@ -6,7 +6,7 @@ import (
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
 )
 
-//go:generate mockery -name FlyteAdminDataProvider -output=../mocks -case=underscore
+//go:generate mockery-v2 --name=FlyteAdminDataProvider --output=../mocks --case=underscore --with-expecter
 
 type FlyteAdminDataProvider interface {
 	GetClusterResourceAttributes(ctx context.Context, project, domain string) (*admin.ClusterResourceAttributes, error)

diff --git a/flyteadmin/pkg/clusterresource/mocks/flyte_admin_data_provider.go b/flyteadmin/pkg/clusterresource/mocks/flyte_admin_data_provider.go