mirror of
https://github.com/k3s-io/kubernetes.git
synced 2025-08-11 13:02:14 +00:00
feature(dra): add mock workqueue for controller unit test
This commit is contained in:
parent
dc481fef90
commit
1c67ed08f2
@ -0,0 +1,165 @@
|
||||
/*
|
||||
Copyright 2024 The Kubernetes Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package controller
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"k8s.io/client-go/util/workqueue"
|
||||
)
|
||||
|
||||
// TODO (pohly): move this to k8s.io/client-go/util/workqueue/workqueue.go
|
||||
// if it turns out to be generally useful. Doc comments are already written
|
||||
// as if the code was there.
|
||||
|
||||
// MockQueue is an implementation of [TypedRateLimitingInterface] which
|
||||
// can be used to test a function which pulls work items out of a queue
|
||||
// and processes them.
|
||||
//
|
||||
// A null instance is directly usable. The usual usage is:
|
||||
//
|
||||
// var m workqueue.Mock[string]
|
||||
// m.SyncOne("some-item", func(queue workqueue.TypedRateLimitingInterface[string]) { ... } )
|
||||
// if diff := cmp.Diff(workqueue.Mock[string]{}, m); diff != "" {
|
||||
// t.Errorf("unexpected state of mock work queue after sync (-want, +got):\n%s", diff)
|
||||
// }
|
||||
//
|
||||
// All slices get reset to nil when they become empty, so there are no spurious
|
||||
// differences because of the nil vs. empty slice.
|
||||
type Mock[T comparable] struct {
|
||||
// Ready contains the items which are ready for processing.
|
||||
Ready []T
|
||||
|
||||
// InFlight contains the items which are currently being processed (= Get
|
||||
// was called, Done not yet).
|
||||
InFlight []T
|
||||
|
||||
// MismatchedDone contains the items for which Done was called without
|
||||
// a matching Get.
|
||||
MismatchedDone []T
|
||||
|
||||
// Later contains the items which are meant to be added to the queue after
|
||||
// a certain delay (= AddAfter was called for them).
|
||||
Later []MockDelayedItem[T]
|
||||
|
||||
// Failures contains the items and their retry count which failed to be
|
||||
// processed (AddRateLimited called at least once, Forget not yet).
|
||||
// The retry count is always larger than zero.
|
||||
Failures map[T]int
|
||||
|
||||
// ShutDownCalled tracks how often ShutDown got called.
|
||||
ShutDownCalled int
|
||||
|
||||
// ShutDownWithDrainCalled tracks how often ShutDownWithDrain got called.
|
||||
ShutDownWithDrainCalled int
|
||||
}
|
||||
|
||||
// MockDelayedItem is an item which was queue for later processing.
|
||||
type MockDelayedItem[T comparable] struct {
|
||||
Item T
|
||||
Duration time.Duration
|
||||
}
|
||||
|
||||
// SyncOne adds the item to the work queue and calls sync.
|
||||
// That sync function can pull one or more items from the work
|
||||
// queue until the queue is empty. Then it is told that the queue
|
||||
// is shutting down, which must cause it to return.
|
||||
//
|
||||
// The test can then retrieve the state of the queue to check the result.
|
||||
func (m *Mock[T]) SyncOne(item T, sync func(workqueue.TypedRateLimitingInterface[T])) {
|
||||
m.Ready = append(m.Ready, item)
|
||||
sync(m)
|
||||
}
|
||||
|
||||
// Add implements [TypedInterface].
|
||||
func (m *Mock[T]) Add(item T) {
|
||||
m.Ready = append(m.Ready, item)
|
||||
}
|
||||
|
||||
// Len implements [TypedInterface].
|
||||
func (m *Mock[T]) Len() int {
|
||||
return len(m.Ready)
|
||||
}
|
||||
|
||||
// Get implements [TypedInterface].
|
||||
func (m *Mock[T]) Get() (item T, shutdown bool) {
|
||||
if len(m.Ready) == 0 {
|
||||
shutdown = true
|
||||
return
|
||||
}
|
||||
item = m.Ready[0]
|
||||
m.Ready = m.Ready[1:]
|
||||
if len(m.Ready) == 0 {
|
||||
m.Ready = nil
|
||||
}
|
||||
m.InFlight = append(m.InFlight, item)
|
||||
return item, false
|
||||
}
|
||||
|
||||
// Done implements [TypedInterface].
|
||||
func (m *Mock[T]) Done(item T) {
|
||||
index := slices.Index(m.InFlight, item)
|
||||
if index < 0 {
|
||||
m.MismatchedDone = append(m.MismatchedDone, item)
|
||||
}
|
||||
m.InFlight = slices.Delete(m.InFlight, index, index+1)
|
||||
if len(m.InFlight) == 0 {
|
||||
m.InFlight = nil
|
||||
}
|
||||
}
|
||||
|
||||
// ShutDown implements [TypedInterface].
|
||||
func (m *Mock[T]) ShutDown() {
|
||||
m.ShutDownCalled++
|
||||
}
|
||||
|
||||
// ShutDownWithDrain implements [TypedInterface].
|
||||
func (m *Mock[T]) ShutDownWithDrain() {
|
||||
m.ShutDownWithDrainCalled++
|
||||
}
|
||||
|
||||
// ShuttingDown implements [TypedInterface].
|
||||
func (m *Mock[T]) ShuttingDown() bool {
|
||||
return m.ShutDownCalled > 0 || m.ShutDownWithDrainCalled > 0
|
||||
}
|
||||
|
||||
// AddAfter implements [TypedDelayingInterface.AddAfter]
|
||||
func (m *Mock[T]) AddAfter(item T, duration time.Duration) {
|
||||
m.Later = append(m.Later, MockDelayedItem[T]{Item: item, Duration: duration})
|
||||
}
|
||||
|
||||
// AddRateLimited implements [TypedRateLimitingInterface.AddRateLimited].
|
||||
func (m *Mock[T]) AddRateLimited(item T) {
|
||||
if m.Failures == nil {
|
||||
m.Failures = make(map[T]int)
|
||||
}
|
||||
m.Failures[item]++
|
||||
}
|
||||
|
||||
// Forget implements [TypedRateLimitingInterface.Forget].
|
||||
func (m *Mock[T]) Forget(item T) {
|
||||
if m.Failures == nil {
|
||||
return
|
||||
}
|
||||
delete(m.Failures, item)
|
||||
}
|
||||
|
||||
// NumRequeues implements [TypedRateLimitingInterface.NumRequeues].
|
||||
func (m *Mock[T]) NumRequeues(item T) int {
|
||||
return m.Failures[item]
|
||||
}
|
Loading…
Reference in New Issue
Block a user