github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/dataprotection/utils/periodical_enqueue_source_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 utils
    21  
    22  import (
    23  	"context"
    24  	"time"
    25  
    26  	. "github.com/onsi/ginkgo/v2"
    27  	. "github.com/onsi/gomega"
    28  
    29  	"k8s.io/apimachinery/pkg/api/meta"
    30  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    31  	"k8s.io/apimachinery/pkg/runtime"
    32  	"k8s.io/client-go/kubernetes/scheme"
    33  	"k8s.io/client-go/util/workqueue"
    34  	"sigs.k8s.io/controller-runtime/pkg/client"
    35  	"sigs.k8s.io/controller-runtime/pkg/client/fake"
    36  	"sigs.k8s.io/controller-runtime/pkg/event"
    37  	"sigs.k8s.io/controller-runtime/pkg/predicate"
    38  	"sigs.k8s.io/controller-runtime/pkg/reconcile"
    39  
    40  	dpv1alpha1 "github.com/1aal/kubeblocks/apis/dataprotection/v1alpha1"
    41  )
    42  
    43  var _ = Describe("Periodical Enqueue Source", func() {
    44  	const (
    45  		backupName = "test-backup"
    46  	)
    47  
    48  	var (
    49  		ctx        context.Context
    50  		cancelFunc context.CancelFunc
    51  		cli        client.Client
    52  		queue      workqueue.RateLimitingInterface
    53  	)
    54  
    55  	createBackup := func(name string) {
    56  		Expect(cli.Create(ctx, &dpv1alpha1.Backup{
    57  			ObjectMeta: metav1.ObjectMeta{
    58  				Name: name,
    59  			},
    60  		})).Should(Succeed())
    61  	}
    62  
    63  	BeforeEach(func() {
    64  		Expect(dpv1alpha1.AddToScheme(scheme.Scheme)).Should(Succeed())
    65  		ctx, cancelFunc = context.WithCancel(context.TODO())
    66  		cli = (&fake.ClientBuilder{}).Build()
    67  		queue = workqueue.NewRateLimitingQueue(workqueue.DefaultControllerRateLimiter())
    68  	})
    69  
    70  	Context("source", func() {
    71  		var source *PeriodicalEnqueueSource
    72  
    73  		BeforeEach(func() {
    74  			By("create source")
    75  			source = NewPeriodicalEnqueueSource(cli, &dpv1alpha1.BackupList{}, 1*time.Second, PeriodicalEnqueueSourceOption{})
    76  		})
    77  
    78  		It("should start success", func() {
    79  			By("start source")
    80  			Expect(source.Start(ctx, nil, queue)).Should(Succeed())
    81  
    82  			By("wait and there is no resources")
    83  			time.Sleep(1 * time.Second)
    84  			Expect(queue.Len()).Should(Equal(0))
    85  
    86  			By("create a resource")
    87  			createBackup(backupName)
    88  
    89  			By("wait and there is one resource")
    90  			time.Sleep(2 * time.Second)
    91  			Expect(queue.Len()).Should(Equal(1))
    92  
    93  			By("cancel context, the queue source shouldn't run anymore")
    94  			item, _ := queue.Get()
    95  			queue.Forget(item)
    96  			Expect(queue.Len()).Should(Equal(0))
    97  			cancelFunc()
    98  			time.Sleep(2 * time.Second)
    99  			Expect(queue.Len()).Should(Equal(0))
   100  		})
   101  
   102  		It("predicate should work", func() {
   103  			By("start source")
   104  			Expect(source.Start(ctx, nil, queue, predicate.Funcs{
   105  				GenericFunc: func(event event.GenericEvent) bool {
   106  					return event.Object.GetName() == backupName
   107  				},
   108  			}))
   109  
   110  			By("create a resource match predicate")
   111  			createBackup(backupName)
   112  
   113  			By("create another resource that does not match predicate")
   114  			createBackup(backupName + "-1")
   115  
   116  			By("wait and there is one resource")
   117  			time.Sleep(2 * time.Second)
   118  			Expect(queue.Len()).Should(Equal(1))
   119  
   120  			cancelFunc()
   121  		})
   122  
   123  		It("order function should work", func() {
   124  			By("set source order func")
   125  			source.option.OrderFunc = func(objList client.ObjectList) client.ObjectList {
   126  				backupList := &dpv1alpha1.BackupList{}
   127  				objArray := make([]runtime.Object, 0)
   128  				backups, _ := meta.ExtractList(objList)
   129  				objArray = append(objArray, backups[1], backups[0])
   130  				_ = meta.SetList(backupList, objArray)
   131  				return backupList
   132  			}
   133  
   134  			By("create a resource")
   135  			createBackup(backupName + "-1")
   136  
   137  			By("create another resource")
   138  			createBackup(backupName + "-2")
   139  
   140  			By("start source")
   141  			Expect(source.Start(ctx, nil, queue)).Should(Succeed())
   142  
   143  			time.Sleep(2 * time.Second)
   144  			Expect(queue.Len()).Should(Equal(2))
   145  			first, _ := queue.Get()
   146  			Expect(first.(reconcile.Request).Name).Should(Equal(backupName + "-2"))
   147  
   148  			cancelFunc()
   149  		})
   150  	})
   151  })