repo_description: The official CLI for Amazon EKS
repo_topics: amazon-web-services aws-cloudformation aws-ec2 aws-eks eks kubernetes kubernetes-cluster kubernetes-deployment kubernetes-setup
code: package drain_test
  
  import (
  	"context"
  	"errors"
  	"fmt"
  	"time"
  
  	"golang.org/x/sync/semaphore"
  	corev1 "k8s.io/api/core/v1"
  	apierrors "k8s.io/apimachinery/pkg/api/errors"
  
  	"github.com/weaveworks/eksctl/pkg/drain/evictor"
  
  	"github.com/weaveworks/eksctl/pkg/drain/fakes"
  
  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
  
  	. "github.com/onsi/ginkgo/v2"
  	. "github.com/onsi/gomega"
  
  	"k8s.io/client-go/kubernetes/fake"
  
  	"github.com/weaveworks/eksctl/pkg/drain"
  	"github.com/weaveworks/eksctl/pkg/eks/mocks"
  )
  
  var _ = Describe("Drain", func() {
  	var (
  		mockNG        mocks.KubeNodeGroup
  		fakeClientSet *fake.Clientset
  		fakeEvictor   *fakes.FakeEvictor
  		nodeName      = "node-1"
  		ctx           context.Context
  		sem           *semaphore.Weighted
  	)
  
  	BeforeEach(func() {
  		mockNG = mocks.KubeNodeGroup{}
  		mockNG.Mock.On("NameString").Return("node-1")
  		mockNG.Mock.On("ListOptions").Return(metav1.ListOptions{})
  		fakeClientSet = fake.NewSimpleClientset()
  		fakeEvictor = new(fakes.FakeEvictor)
  		var cancel context.CancelFunc
  		ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
  		DeferCleanup(cancel)
  		sem = semaphore.NewWeighted(1)
  	})
  
  	When("all nodes drain successfully", func() {
  		var pod corev1.Pod
  
  		BeforeEach(func() {
  			pod = corev1.Pod{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: "pod-1",
  				},
  			}
  
  			fakeEvictor.GetPodsForEvictionReturnsOnCall(0, &evictor.PodDeleteList{
  				Items: []evictor.PodDelete{
  					{
  						Pod: pod,
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  				},
  			}, nil)
  			fakeEvictor.GetPodsForEvictionReturnsOnCall(1, &evictor.PodDeleteList{}, nil)
  
  			fakeEvictor.EvictOrDeletePodReturns(nil)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: false,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("does not error", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second*10, time.Second, 0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			err := nodeGroupDrainer.Drain(ctx, sem)
  			Expect(err).NotTo(HaveOccurred())
  
  			Expect(fakeEvictor.GetPodsForEvictionCallCount()).To(BeNumerically(">=", 2))
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(Equal(1))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(0)).To(Equal(pod))
  
  			node, err := fakeClientSet.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
  			Expect(err).NotTo(HaveOccurred())
  			Expect(node.Spec.Unschedulable).To(BeTrue())
  		})
  	})
  
  	When("the nodes never drain successfully", func() {
  		var pod corev1.Pod
  
  		BeforeEach(func() {
  			pod = corev1.Pod{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: "pod-1",
  				},
  			}
  
  			fakeEvictor.GetPodsForEvictionReturns(&evictor.PodDeleteList{
  				Items: []evictor.PodDelete{
  					{
  						Pod: pod,
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  				},
  			}, nil)
  
  			fakeEvictor.EvictOrDeletePodReturns(nil)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("times out and errors", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, 0, time.Second, 0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
  			defer cancel()
  			err := nodeGroupDrainer.Drain(ctx, sem)
  			Expect(err).To(MatchError(`timed out waiting for nodegroup "node-1" to be drained`))
  		})
  	})
  
  	When("Evictions are not supported", func() {
  
  		BeforeEach(func() {
  			fakeEvictor.CanUseEvictionsReturns(fmt.Errorf("error1"))
  		})
  
  		It("errors", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, 0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
  			defer cancel()
  			err := nodeGroupDrainer.Drain(ctx, sem)
  			Expect(err).To(MatchError("checking if cluster implements policy API: error1"))
  		})
  	})
  
  	When("Evictions are disabled", func() {
  		var pod corev1.Pod
  
  		BeforeEach(func() {
  			pod = corev1.Pod{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: "pod-1",
  				},
  			}
  
  			fakeEvictor.GetPodsForEvictionReturnsOnCall(0, &evictor.PodDeleteList{
  				Items: []evictor.PodDelete{
  					{
  						Pod: pod,
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  				},
  			}, nil)
  			fakeEvictor.GetPodsForEvictionReturnsOnCall(1, &evictor.PodDeleteList{}, nil)
  
  			fakeEvictor.EvictOrDeletePodReturns(nil)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: false,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("does not error", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, time.Second*0, false, true, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			Expect(nodeGroupDrainer.Drain(ctx, sem)).To(Succeed())
  
  			Expect(fakeEvictor.GetPodsForEvictionCallCount()).To(BeNumerically(">=", 2))
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(Equal(1))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(0)).To(Equal(pod))
  
  			node, err := fakeClientSet.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
  			Expect(err).NotTo(HaveOccurred())
  			Expect(node.Spec.Unschedulable).To(BeTrue())
  		})
  	})
  
  	When("undo is true", func() {
  
  		BeforeEach(func() {
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: true,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("uncordons all the nodes", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, time.Second*0, true, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			err := nodeGroupDrainer.Drain(ctx, sem)
  			Expect(err).NotTo(HaveOccurred())
  
  			node, err := fakeClientSet.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
  			Expect(err).NotTo(HaveOccurred())
  			Expect(node.Spec.Unschedulable).To(BeFalse())
  
  			Expect(fakeEvictor.GetPodsForEvictionCallCount()).To(BeZero())
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(BeZero())
  		})
  	})
  
  	When("an eviction fails recoverably", func() {
  		var pod corev1.Pod
  
  		BeforeEach(func() {
  			pod = corev1.Pod{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: "pod-1",
  				},
  			}
  
  			for i := 0; i < 2; i++ {
  				fakeEvictor.GetPodsForEvictionReturnsOnCall(i, &evictor.PodDeleteList{
  					Items: []evictor.PodDelete{
  						{
  							Pod: pod,
  							Status: evictor.PodDeleteStatus{
  								Delete: true,
  							},
  						},
  					},
  				}, nil)
  			}
  			fakeEvictor.GetPodsForEvictionReturnsOnCall(2, nil, nil)
  
  			fakeEvictor.EvictOrDeletePodReturnsOnCall(0, apierrors.NewTooManyRequestsError("error1"))
  			fakeEvictor.EvictOrDeletePodReturnsOnCall(1, nil)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: false,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("does not error", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, 0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			Expect(nodeGroupDrainer.Drain(ctx, sem)).To(Succeed())
  
  			Expect(fakeEvictor.GetPodsForEvictionCallCount()).To(Equal(3))
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(Equal(2))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(0)).To(Equal(pod))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(1)).To(Equal(pod))
  		})
  	})
  
  	When("an eviction fails irrecoverably", func() {
  		var pod corev1.Pod
  		var evictionError error
  
  		BeforeEach(func() {
  			pod = corev1.Pod{
  				ObjectMeta: metav1.ObjectMeta{
  					Name:      "pod-1",
  					Namespace: "ns-1",
  				},
  			}
  
  			fakeEvictor.GetPodsForEvictionReturns(&evictor.PodDeleteList{
  				Items: []evictor.PodDelete{
  					{
  						Pod: pod,
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  				},
  			}, nil)
  
  			evictionError = errors.New("error1")
  			fakeEvictor.EvictOrDeletePodReturns(evictionError)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: false,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("returns an error", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, 0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			err := nodeGroupDrainer.Drain(ctx, sem)
  			Expect(err.Error()).To(ContainSubstring("unrecoverable error evicting pod: ns-1/pod-1"))
  
  			Expect(fakeEvictor.GetPodsForEvictionCallCount()).To(Equal(1))
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(Equal(1))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(0)).To(Equal(pod))
  		})
  	})
  
  	When("eviction fails recoverably with multiple pods", func() {
  		var pods []corev1.Pod
  		var evictionError error
  
  		BeforeEach(func() {
  			pods = []corev1.Pod{
  				{
  					ObjectMeta: metav1.ObjectMeta{
  						Name:      "pod-1",
  						Namespace: "ns-1",
  					},
  				},
  				{
  					ObjectMeta: metav1.ObjectMeta{
  						Name:      "pod-2",
  						Namespace: "ns-2",
  					},
  				},
  			}
  
  			fakeEvictor.GetPodsForEvictionReturns(&evictor.PodDeleteList{
  				Items: []evictor.PodDelete{
  					{
  						Pod: pods[0],
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  					{
  						Pod: pods[1],
  						Status: evictor.PodDeleteStatus{
  							Delete: true,
  						},
  					},
  				},
  			}, nil)
  
  			evictionError = apierrors.NewTooManyRequestsError("error1")
  			fakeEvictor.EvictOrDeletePodReturns(evictionError)
  
  			_, err := fakeClientSet.CoreV1().Nodes().Create(context.Background(), &corev1.Node{
  				ObjectMeta: metav1.ObjectMeta{
  					Name: nodeName,
  				},
  				Spec: corev1.NodeSpec{
  					Unschedulable: false,
  				},
  			}, metav1.CreateOptions{})
  			Expect(err).NotTo(HaveOccurred())
  		})
  
  		It("it attempts to drain all pods", func() {
  			nodeGroupDrainer := drain.NewNodeGroupDrainer(fakeClientSet, &mockNG, time.Second, time.Second, time.Second*0, false, false, 1)
  			nodeGroupDrainer.SetDrainer(fakeEvictor)
  
  			_ = nodeGroupDrainer.Drain(ctx, sem)
  
  			Expect(fakeEvictor.EvictOrDeletePodCallCount()).To(BeNumerically(">=", 2))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(0)).To(Equal(pods[0]))
  			Expect(fakeEvictor.EvictOrDeletePodArgsForCall(1)).To(Equal(pods[1]))
  		})
  	})
  })
  