From 737bf16043f3d32c25148c56094197449bacdfe0 Mon Sep 17 00:00:00 2001 From: Illia Litovchenko Date: Thu, 19 Feb 2026 10:12:40 +0000 Subject: [PATCH] Mocks for k8s client/node drainer --- internal/k8s/kubernetes.go | 43 +++++--- internal/k8s/mock/kubernetes.go | 170 ++++++++++++++++++++++++++++++ internal/nodes/mock/kubernetes.go | 67 ++++++++++++ internal/nodes/node_drainer.go | 14 +-- 4 files changed, 276 insertions(+), 18 deletions(-) create mode 100644 internal/k8s/mock/kubernetes.go create mode 100644 internal/nodes/mock/kubernetes.go diff --git a/internal/k8s/kubernetes.go b/internal/k8s/kubernetes.go index d677bbe9..b8d5ac65 100644 --- a/internal/k8s/kubernetes.go +++ b/internal/k8s/kubernetes.go @@ -1,3 +1,5 @@ +//go:generate mockgen -destination ./mock/kubernetes.go . Client + package k8s import ( @@ -44,32 +46,49 @@ const ( DefaultMaxRetriesK8SOperation = 5 ) +type Client interface { + PatchNode(ctx context.Context, node *v1.Node, changeFn func(*v1.Node)) error + PatchNodeStatus(ctx context.Context, name string, patch []byte) error + EvictPod(ctx context.Context, pod v1.Pod, podEvictRetryDelay time.Duration, version schema.GroupVersion) error + CordonNode(ctx context.Context, node *v1.Node) error + GetNodeByIDs(ctx context.Context, nodeName, nodeID, providerID string) (*v1.Node, error) + ExecuteBatchPodActions( + ctx context.Context, + pods []*v1.Pod, + action func(context.Context, v1.Pod) error, + actionName string, + ) ([]*v1.Pod, []PodActionFailure) + DeletePod(ctx context.Context, options metav1.DeleteOptions, pod v1.Pod, podDeleteRetries int, podDeleteRetryDelay time.Duration) error + Clientset() kubernetes.Interface + Log() logrus.FieldLogger +} + // Client provides Kubernetes operations with common dependencies. -type Client struct { +type client struct { clientset kubernetes.Interface log logrus.FieldLogger } // NewClient creates a new K8s client with the given dependencies. -func NewClient(clientset kubernetes.Interface, log logrus.FieldLogger) *Client { - return &Client{ +func NewClient(clientset kubernetes.Interface, log logrus.FieldLogger) Client { + return &client{ clientset: clientset, log: log, } } // Clientset returns the underlying kubernetes.Interface. -func (c *Client) Clientset() kubernetes.Interface { +func (c *client) Clientset() kubernetes.Interface { return c.clientset } // Log returns the logger. -func (c *Client) Log() logrus.FieldLogger { +func (c *client) Log() logrus.FieldLogger { return c.log } // PatchNode patches a node with the given change function. -func (c *Client) PatchNode(ctx context.Context, node *v1.Node, changeFn func(*v1.Node)) error { +func (c *client) PatchNode(ctx context.Context, node *v1.Node, changeFn func(*v1.Node)) error { logger := logger.FromContext(ctx, c.log) oldData, err := json.Marshal(node) if err != nil { @@ -108,7 +127,7 @@ func (c *Client) PatchNode(ctx context.Context, node *v1.Node, changeFn func(*v1 } // PatchNodeStatus patches the status of a node. -func (c *Client) PatchNodeStatus(ctx context.Context, name string, patch []byte) error { +func (c *client) PatchNodeStatus(ctx context.Context, name string, patch []byte) error { logger := logger.FromContext(ctx, c.log) err := waitext.Retry( @@ -134,7 +153,7 @@ func (c *Client) PatchNodeStatus(ctx context.Context, name string, patch []byte) return nil } -func (c *Client) CordonNode(ctx context.Context, node *v1.Node) error { +func (c *client) CordonNode(ctx context.Context, node *v1.Node) error { if node.Spec.Unschedulable { return nil } @@ -149,7 +168,7 @@ func (c *Client) CordonNode(ctx context.Context, node *v1.Node) error { } // GetNodeByIDs retrieves a node by name and validates its ID and provider ID. -func (c *Client) GetNodeByIDs(ctx context.Context, nodeName, nodeID, providerID string) (*v1.Node, error) { +func (c *client) GetNodeByIDs(ctx context.Context, nodeName, nodeID, providerID string) (*v1.Node, error) { if nodeID == "" && providerID == "" { return nil, fmt.Errorf("node and provider IDs are empty %w", ErrAction) } @@ -178,7 +197,7 @@ func (c *Client) GetNodeByIDs(ctx context.Context, nodeName, nodeID, providerID // It does internal throttling to avoid spawning a goroutine-per-pod on large lists. // Returns two sets of pods - the ones that successfully executed the action and the ones that failed. // actionName might be used to distinguish what is the operation (for logs, debugging, etc.) but is optional. -func (c *Client) ExecuteBatchPodActions( +func (c *client) ExecuteBatchPodActions( ctx context.Context, pods []*v1.Pod, action func(context.Context, v1.Pod) error, @@ -250,7 +269,7 @@ func (c *Client) ExecuteBatchPodActions( // EvictPod evicts a pod from a k8s node. Error handling is based on eviction api documentation: // https://kubernetes.io/docs/tasks/administer-cluster/safely-drain-node/#the-eviction-api -func (c *Client) EvictPod(ctx context.Context, pod v1.Pod, podEvictRetryDelay time.Duration, version schema.GroupVersion) error { +func (c *client) EvictPod(ctx context.Context, pod v1.Pod, podEvictRetryDelay time.Duration, version schema.GroupVersion) error { logger := logger.FromContext(ctx, c.log) b := waitext.NewConstantBackoff(podEvictRetryDelay) @@ -306,7 +325,7 @@ func (c *Client) EvictPod(ctx context.Context, pod v1.Pod, podEvictRetryDelay ti } // DeletePod deletes a pod from the cluster. -func (c *Client) DeletePod(ctx context.Context, options metav1.DeleteOptions, pod v1.Pod, podDeleteRetries int, podDeleteRetryDelay time.Duration) error { +func (c *client) DeletePod(ctx context.Context, options metav1.DeleteOptions, pod v1.Pod, podDeleteRetries int, podDeleteRetryDelay time.Duration) error { logger := logger.FromContext(ctx, c.log) b := waitext.NewConstantBackoff(podDeleteRetryDelay) diff --git a/internal/k8s/mock/kubernetes.go b/internal/k8s/mock/kubernetes.go new file mode 100644 index 00000000..a5c2f3dd --- /dev/null +++ b/internal/k8s/mock/kubernetes.go @@ -0,0 +1,170 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/castai/cluster-controller/internal/k8s (interfaces: Client) + +// Package mock_k8s is a generated GoMock package. +package mock_k8s + +import ( + context "context" + reflect "reflect" + time "time" + + k8s "github.com/castai/cluster-controller/internal/k8s" + gomock "github.com/golang/mock/gomock" + logrus "github.com/sirupsen/logrus" + v1 "k8s.io/api/core/v1" + v10 "k8s.io/apimachinery/pkg/apis/meta/v1" + schema "k8s.io/apimachinery/pkg/runtime/schema" + kubernetes "k8s.io/client-go/kubernetes" +) + +// MockClient is a mock of Client interface. +type MockClient struct { + ctrl *gomock.Controller + recorder *MockClientMockRecorder +} + +// MockClientMockRecorder is the mock recorder for MockClient. +type MockClientMockRecorder struct { + mock *MockClient +} + +// NewMockClient creates a new mock instance. +func NewMockClient(ctrl *gomock.Controller) *MockClient { + mock := &MockClient{ctrl: ctrl} + mock.recorder = &MockClientMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockClient) EXPECT() *MockClientMockRecorder { + return m.recorder +} + +// Clientset mocks base method. +func (m *MockClient) Clientset() kubernetes.Interface { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Clientset") + ret0, _ := ret[0].(kubernetes.Interface) + return ret0 +} + +// Clientset indicates an expected call of Clientset. +func (mr *MockClientMockRecorder) Clientset() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Clientset", reflect.TypeOf((*MockClient)(nil).Clientset)) +} + +// CordonNode mocks base method. +func (m *MockClient) CordonNode(arg0 context.Context, arg1 *v1.Node) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CordonNode", arg0, arg1) + ret0, _ := ret[0].(error) + return ret0 +} + +// CordonNode indicates an expected call of CordonNode. +func (mr *MockClientMockRecorder) CordonNode(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CordonNode", reflect.TypeOf((*MockClient)(nil).CordonNode), arg0, arg1) +} + +// DeletePod mocks base method. +func (m *MockClient) DeletePod(arg0 context.Context, arg1 v10.DeleteOptions, arg2 v1.Pod, arg3 int, arg4 time.Duration) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeletePod", arg0, arg1, arg2, arg3, arg4) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeletePod indicates an expected call of DeletePod. +func (mr *MockClientMockRecorder) DeletePod(arg0, arg1, arg2, arg3, arg4 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeletePod", reflect.TypeOf((*MockClient)(nil).DeletePod), arg0, arg1, arg2, arg3, arg4) +} + +// EvictPod mocks base method. +func (m *MockClient) EvictPod(arg0 context.Context, arg1 v1.Pod, arg2 time.Duration, arg3 schema.GroupVersion) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "EvictPod", arg0, arg1, arg2, arg3) + ret0, _ := ret[0].(error) + return ret0 +} + +// EvictPod indicates an expected call of EvictPod. +func (mr *MockClientMockRecorder) EvictPod(arg0, arg1, arg2, arg3 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EvictPod", reflect.TypeOf((*MockClient)(nil).EvictPod), arg0, arg1, arg2, arg3) +} + +// ExecuteBatchPodActions mocks base method. +func (m *MockClient) ExecuteBatchPodActions(arg0 context.Context, arg1 []*v1.Pod, arg2 func(context.Context, v1.Pod) error, arg3 string) ([]*v1.Pod, []k8s.PodActionFailure) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ExecuteBatchPodActions", arg0, arg1, arg2, arg3) + ret0, _ := ret[0].([]*v1.Pod) + ret1, _ := ret[1].([]k8s.PodActionFailure) + return ret0, ret1 +} + +// ExecuteBatchPodActions indicates an expected call of ExecuteBatchPodActions. +func (mr *MockClientMockRecorder) ExecuteBatchPodActions(arg0, arg1, arg2, arg3 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExecuteBatchPodActions", reflect.TypeOf((*MockClient)(nil).ExecuteBatchPodActions), arg0, arg1, arg2, arg3) +} + +// GetNodeByIDs mocks base method. +func (m *MockClient) GetNodeByIDs(arg0 context.Context, arg1, arg2, arg3 string) (*v1.Node, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetNodeByIDs", arg0, arg1, arg2, arg3) + ret0, _ := ret[0].(*v1.Node) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetNodeByIDs indicates an expected call of GetNodeByIDs. +func (mr *MockClientMockRecorder) GetNodeByIDs(arg0, arg1, arg2, arg3 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetNodeByIDs", reflect.TypeOf((*MockClient)(nil).GetNodeByIDs), arg0, arg1, arg2, arg3) +} + +// Log mocks base method. +func (m *MockClient) Log() logrus.FieldLogger { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Log") + ret0, _ := ret[0].(logrus.FieldLogger) + return ret0 +} + +// Log indicates an expected call of Log. +func (mr *MockClientMockRecorder) Log() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Log", reflect.TypeOf((*MockClient)(nil).Log)) +} + +// PatchNode mocks base method. +func (m *MockClient) PatchNode(arg0 context.Context, arg1 *v1.Node, arg2 func(*v1.Node)) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PatchNode", arg0, arg1, arg2) + ret0, _ := ret[0].(error) + return ret0 +} + +// PatchNode indicates an expected call of PatchNode. +func (mr *MockClientMockRecorder) PatchNode(arg0, arg1, arg2 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PatchNode", reflect.TypeOf((*MockClient)(nil).PatchNode), arg0, arg1, arg2) +} + +// PatchNodeStatus mocks base method. +func (m *MockClient) PatchNodeStatus(arg0 context.Context, arg1 string, arg2 []byte) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PatchNodeStatus", arg0, arg1, arg2) + ret0, _ := ret[0].(error) + return ret0 +} + +// PatchNodeStatus indicates an expected call of PatchNodeStatus. +func (mr *MockClientMockRecorder) PatchNodeStatus(arg0, arg1, arg2 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PatchNodeStatus", reflect.TypeOf((*MockClient)(nil).PatchNodeStatus), arg0, arg1, arg2) +} diff --git a/internal/nodes/mock/kubernetes.go b/internal/nodes/mock/kubernetes.go new file mode 100644 index 00000000..5a739075 --- /dev/null +++ b/internal/nodes/mock/kubernetes.go @@ -0,0 +1,67 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/castai/cluster-controller/internal/nodes (interfaces: Drainer) + +// Package mock_nodes is a generated GoMock package. +package mock_nodes + +import ( + context "context" + reflect "reflect" + + nodes "github.com/castai/cluster-controller/internal/nodes" + gomock "github.com/golang/mock/gomock" + v1 "k8s.io/api/core/v1" +) + +// MockDrainer is a mock of Drainer interface. +type MockDrainer struct { + ctrl *gomock.Controller + recorder *MockDrainerMockRecorder +} + +// MockDrainerMockRecorder is the mock recorder for MockDrainer. +type MockDrainerMockRecorder struct { + mock *MockDrainer +} + +// NewMockDrainer creates a new mock instance. +func NewMockDrainer(ctrl *gomock.Controller) *MockDrainer { + mock := &MockDrainer{ctrl: ctrl} + mock.recorder = &MockDrainerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDrainer) EXPECT() *MockDrainerMockRecorder { + return m.recorder +} + +// Drain mocks base method. +func (m *MockDrainer) Drain(arg0 context.Context, arg1 nodes.DrainRequest) ([]*v1.Pod, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Drain", arg0, arg1) + ret0, _ := ret[0].([]*v1.Pod) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Drain indicates an expected call of Drain. +func (mr *MockDrainerMockRecorder) Drain(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Drain", reflect.TypeOf((*MockDrainer)(nil).Drain), arg0, arg1) +} + +// Evict mocks base method. +func (m *MockDrainer) Evict(arg0 context.Context, arg1 nodes.EvictRequest) ([]*v1.Pod, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Evict", arg0, arg1) + ret0, _ := ret[0].([]*v1.Pod) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Evict indicates an expected call of Evict. +func (mr *MockDrainerMockRecorder) Evict(arg0, arg1 interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Evict", reflect.TypeOf((*MockDrainer)(nil).Evict), arg0, arg1) +} diff --git a/internal/nodes/node_drainer.go b/internal/nodes/node_drainer.go index 55b888f4..2bf4029a 100644 --- a/internal/nodes/node_drainer.go +++ b/internal/nodes/node_drainer.go @@ -1,3 +1,5 @@ +//go:generate mockgen -destination ./mock/kubernetes.go . Drainer + package nodes import ( @@ -44,14 +46,14 @@ type DrainerConfig struct { type drainer struct { pods informer.PodInformer - client *k8s.Client + client k8s.Client cfg DrainerConfig log logrus.FieldLogger } func NewDrainer( pods informer.PodInformer, - client *k8s.Client, + client k8s.Client, log logrus.FieldLogger, cfg DrainerConfig, ) Drainer { @@ -75,7 +77,7 @@ func (d *drainer) Drain(ctx context.Context, data DrainRequest) ([]*core.Pod, er toEvict := d.prioritizePods(pods, data.CastNamespace, data.SkipDeletedTimeoutSeconds) if len(toEvict) == 0 { - return []*core.Pod{}, nil + return nil, nil } _, failed, err := d.tryDrain(ctx, toEvict, data.DeleteOptions) @@ -85,7 +87,7 @@ func (d *drainer) Drain(ctx context.Context, data DrainRequest) ([]*core.Pod, er err = d.waitTerminaition(ctx, data.Node, failed) if err != nil { - return []*core.Pod{}, err + return nil, err } logger.Info("drain finished") @@ -115,7 +117,7 @@ func (d *drainer) Evict(ctx context.Context, data EvictRequest) ([]*core.Pod, er toEvict := d.prioritizePods(pods, data.CastNamespace, data.SkipDeletedTimeoutSeconds) if len(toEvict) == 0 { - return []*core.Pod{}, nil + return nil, nil } _, ignored, err := d.tryEvict(ctx, toEvict) @@ -125,7 +127,7 @@ func (d *drainer) Evict(ctx context.Context, data EvictRequest) ([]*core.Pod, er err = d.waitTerminaition(ctx, data.Node, ignored) if err != nil { - return []*core.Pod{}, err + return nil, err } logger.Info("eviction finished")