@@ -18,11 +18,12 @@ import (
1818 "context"
1919 "testing"
2020
21+ "github.com/golang/mock/gomock"
2122 "github.com/stretchr/testify/assert"
22- "github.com/stretchr/testify/mock"
2323 corev1 "k8s.io/api/core/v1"
2424 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2525 "k8s.io/apimachinery/pkg/runtime"
26+ "sigs.k8s.io/controller-runtime/pkg/client"
2627
2728 "github.com/banzaicloud/koperator/pkg/scale"
2829
@@ -370,8 +371,11 @@ func TestCreateCCOperation(t *testing.T) {
370371 },
371372 }
372373
374+ mockCtrl := gomock .NewController (t )
375+
373376 for _ , testCase := range testCases {
374- mockClient := new (mocks.Client )
377+ mockClient := mocks .NewMockClient (mockCtrl )
378+ mockSubResourceClient := mocks .NewMockSubResourceClient (mockCtrl )
375379 scheme := runtime .NewScheme ()
376380 _ = v1beta1 .AddToScheme (scheme )
377381 _ = corev1 .AddToScheme (scheme )
@@ -389,17 +393,17 @@ func TestCreateCCOperation(t *testing.T) {
389393
390394 // Mock the Create call and capture the operation
391395 var createdOperation * banzaiv1alpha1.CruiseControlOperation
392- mockClient .On ( " Create" , ctx , mock . IsType (& banzaiv1alpha1.CruiseControlOperation {})).Run (func (args mock. Arguments ) {
393- createdOperation = args . Get ( 1 ) .(* banzaiv1alpha1.CruiseControlOperation )
396+ mockClient .EXPECT (). Create ( ctx , gomock . AssignableToTypeOf (& banzaiv1alpha1.CruiseControlOperation {})).Do (func (ctx context. Context , obj client. Object , opts ... client. CreateOption ) {
397+ createdOperation = obj .(* banzaiv1alpha1.CruiseControlOperation )
394398 createdOperation .ObjectMeta .Name = "generated-name"
395399 }).Return (nil )
396400
397401 // Mock the Status call
398- mockClient .On ( " Status" ).Return (mockClient )
402+ mockClient .EXPECT (). Status ( ).Return (mockSubResourceClient )
399403
400404 // Mock the Update call
401- mockClient . On ( " Update" , ctx , mock . IsType (& banzaiv1alpha1.CruiseControlOperation {})).Run (func (args mock. Arguments ) {
402- arg := args . Get ( 1 ) .(* banzaiv1alpha1.CruiseControlOperation )
405+ mockSubResourceClient . EXPECT (). Update ( ctx , gomock . AssignableToTypeOf (& banzaiv1alpha1.CruiseControlOperation {})).Do (func (ctx context. Context , obj client. Object , opts ... client. UpdateOption ) {
406+ arg := obj .(* banzaiv1alpha1.CruiseControlOperation )
403407 createdOperation .Status = arg .Status
404408 }).Return (nil )
405409
0 commit comments