Fix unit tests after rebase

controllers/cruisecontroltask_controller_test.go

@@ -18,11 +18,12 @@ import (
 	"context"
 	"testing"
 
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"github.com/banzaicloud/koperator/pkg/scale"
 
@@ -370,8 +371,11 @@ func TestCreateCCOperation(t *testing.T) {
 		},
 	}
 
+	mockCtrl := gomock.NewController(t)
+
 	for _, testCase := range testCases {
-		mockClient := new(mocks.Client)
+		mockClient := mocks.NewMockClient(mockCtrl)
+		mockSubResourceClient := mocks.NewMockSubResourceClient(mockCtrl)
 		scheme := runtime.NewScheme()
 		_ = v1beta1.AddToScheme(scheme)
 		_ = corev1.AddToScheme(scheme)
@@ -389,17 +393,17 @@ func TestCreateCCOperation(t *testing.T) {
 
 		// Mock the Create call and capture the operation
 		var createdOperation *banzaiv1alpha1.CruiseControlOperation
-		mockClient.On("Create", ctx, mock.IsType(&banzaiv1alpha1.CruiseControlOperation{})).Run(func(args mock.Arguments) {
-			createdOperation = args.Get(1).(*banzaiv1alpha1.CruiseControlOperation)
+		mockClient.EXPECT().Create(ctx, gomock.AssignableToTypeOf(&banzaiv1alpha1.CruiseControlOperation{})).Do(func(ctx context.Context, obj client.Object, opts ...client.CreateOption) {
+			createdOperation = obj.(*banzaiv1alpha1.CruiseControlOperation)
 			createdOperation.ObjectMeta.Name = "generated-name"
 		}).Return(nil)
 
 		// Mock the Status call
-		mockClient.On("Status").Return(mockClient)
+		mockClient.EXPECT().Status().Return(mockSubResourceClient)
 
 		// Mock the Update call
-		mockClient.On("Update", ctx, mock.IsType(&banzaiv1alpha1.CruiseControlOperation{})).Run(func(args mock.Arguments) {
-			arg := args.Get(1).(*banzaiv1alpha1.CruiseControlOperation)
+		mockSubResourceClient.EXPECT().Update(ctx, gomock.AssignableToTypeOf(&banzaiv1alpha1.CruiseControlOperation{})).Do(func(ctx context.Context, obj client.Object, opts ...client.UpdateOption) {
+			arg := obj.(*banzaiv1alpha1.CruiseControlOperation)
 			createdOperation.Status = arg.Status
 		}).Return(nil)
 

go.mod

@@ -38,6 +38,7 @@ require (
 require (
 	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	golang.org/x/tools v0.7.0 // indirect
 )