github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/controller/rsm/update_plan_test.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package rsm
    21  
    22  import (
    23  	. "github.com/onsi/ginkgo/v2"
    24  	. "github.com/onsi/gomega"
    25  
    26  	apps "k8s.io/api/apps/v1"
    27  	corev1 "k8s.io/api/core/v1"
    28  	"k8s.io/apimachinery/pkg/util/sets"
    29  
    30  	workloads "github.com/1aal/kubeblocks/apis/workloads/v1alpha1"
    31  	"github.com/1aal/kubeblocks/pkg/controller/builder"
    32  )
    33  
    34  var _ = Describe("update plan test.", func() {
    35  	BeforeEach(func() {
    36  		rsm = builder.NewReplicatedStateMachineBuilder(namespace, name).SetRoles(roles).GetObject()
    37  		rsm.Status.UpdateRevision = newRevision
    38  	})
    39  
    40  	Context("plan build&execute", func() {
    41  		var pod0, pod1, pod2, pod3, pod4, pod5, pod6 *corev1.Pod
    42  
    43  		resetPods := func() {
    44  			pod0 = builder.NewPodBuilder(namespace, getPodName(name, 0)).
    45  				AddLabels(roleLabelKey, "follower").
    46  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    47  				GetObject()
    48  
    49  			pod1 = builder.NewPodBuilder(namespace, getPodName(name, 1)).
    50  				AddLabels(roleLabelKey, "logger").
    51  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    52  				GetObject()
    53  
    54  			pod2 = builder.NewPodBuilder(namespace, getPodName(name, 2)).
    55  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    56  				GetObject()
    57  
    58  			pod3 = builder.NewPodBuilder(namespace, getPodName(name, 3)).
    59  				AddLabels(roleLabelKey, "learner").
    60  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    61  				GetObject()
    62  
    63  			pod4 = builder.NewPodBuilder(namespace, getPodName(name, 4)).
    64  				AddLabels(roleLabelKey, "candidate").
    65  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    66  				GetObject()
    67  
    68  			pod5 = builder.NewPodBuilder(namespace, getPodName(name, 5)).
    69  				AddLabels(roleLabelKey, "leader").
    70  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    71  				GetObject()
    72  
    73  			pod6 = builder.NewPodBuilder(namespace, getPodName(name, 6)).
    74  				AddLabels(roleLabelKey, "learner").
    75  				AddLabels(apps.StatefulSetRevisionLabel, oldRevision).
    76  				GetObject()
    77  		}
    78  
    79  		buildPodList := func() []corev1.Pod {
    80  			return []corev1.Pod{*pod0, *pod1, *pod2, *pod3, *pod4, *pod5, *pod6}
    81  		}
    82  
    83  		toPodList := func(pods []*corev1.Pod) []corev1.Pod {
    84  			var list []corev1.Pod
    85  			for _, pod := range pods {
    86  				list = append(list, *pod)
    87  			}
    88  			return list
    89  		}
    90  
    91  		equalPodList := func(podList1, podList2 []corev1.Pod) bool {
    92  			set1 := sets.New[string]()
    93  			set2 := sets.New[string]()
    94  			for _, pod := range podList1 {
    95  				set1.Insert(pod.Name)
    96  			}
    97  			for _, pod := range podList2 {
    98  				set2.Insert(pod.Name)
    99  			}
   100  			return set1.Equal(set2)
   101  		}
   102  
   103  		checkPlan := func(expectedPlan [][]*corev1.Pod) {
   104  			for i, expectedPods := range expectedPlan {
   105  				if i > 0 {
   106  					makePodUpdateReady(newRevision, expectedPlan[i-1]...)
   107  				}
   108  				pods := buildPodList()
   109  				plan := newUpdatePlan(*rsm, pods)
   110  				podUpdateList, err := plan.execute()
   111  				Expect(err).Should(BeNil())
   112  				podList := toPodList(podUpdateList)
   113  				expectedPodList := toPodList(expectedPods)
   114  				Expect(equalPodList(podList, expectedPodList)).Should(BeTrue())
   115  			}
   116  		}
   117  
   118  		BeforeEach(func() {
   119  			resetPods()
   120  		})
   121  
   122  		It("should work well in a serial plan", func() {
   123  			By("build a serial plan")
   124  			strategy := workloads.SerialUpdateStrategy
   125  			rsm.Spec.MemberUpdateStrategy = &strategy
   126  			expectedPlan := [][]*corev1.Pod{
   127  				{pod4},
   128  				{pod2},
   129  				{pod3},
   130  				{pod6},
   131  				{pod1},
   132  				{pod0},
   133  				{pod5},
   134  			}
   135  			checkPlan(expectedPlan)
   136  		})
   137  
   138  		It("should work well in a parallel plan", func() {
   139  			By("build a parallel plan")
   140  			strategy := workloads.ParallelUpdateStrategy
   141  			rsm.Spec.MemberUpdateStrategy = &strategy
   142  			expectedPlan := [][]*corev1.Pod{
   143  				{pod0, pod1, pod2, pod3, pod4, pod5, pod6},
   144  			}
   145  			checkPlan(expectedPlan)
   146  		})
   147  
   148  		It("should work well in a best effort parallel", func() {
   149  			By("build a best effort parallel plan")
   150  			strategy := workloads.BestEffortParallelUpdateStrategy
   151  			rsm.Spec.MemberUpdateStrategy = &strategy
   152  			expectedPlan := [][]*corev1.Pod{
   153  				{pod2, pod3, pod4, pod6},
   154  				{pod1},
   155  				{pod0},
   156  				{pod5},
   157  			}
   158  			checkPlan(expectedPlan)
   159  		})
   160  	})
   161  })