sigs.k8s.io/kueue@v0.6.2/pkg/util/testing/client.go (about)

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package testing
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"strings"
    23  	"sync"
    24  
    25  	"k8s.io/apimachinery/pkg/runtime"
    26  	"k8s.io/apimachinery/pkg/types"
    27  	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
    28  	"sigs.k8s.io/controller-runtime/pkg/client"
    29  	"sigs.k8s.io/controller-runtime/pkg/client/fake"
    30  
    31  	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
    32  	"sigs.k8s.io/kueue/pkg/controller/core/indexer"
    33  )
    34  
    35  func NewFakeClient(objs ...client.Object) client.Client {
    36  	return NewClientBuilder().WithObjects(objs...).WithStatusSubresource(objs...).Build()
    37  }
    38  
    39  func NewClientBuilder(addToSchemes ...func(s *runtime.Scheme) error) *fake.ClientBuilder {
    40  	scheme := runtime.NewScheme()
    41  	if err := clientgoscheme.AddToScheme(scheme); err != nil {
    42  		panic(err)
    43  	}
    44  	if err := kueue.AddToScheme(scheme); err != nil {
    45  		panic(err)
    46  	}
    47  	for i := range addToSchemes {
    48  		if err := addToSchemes[i](scheme); err != nil {
    49  			panic(err)
    50  		}
    51  	}
    52  
    53  	return fake.NewClientBuilder().WithScheme(scheme).
    54  		WithIndex(&kueue.LocalQueue{}, indexer.QueueClusterQueueKey, indexer.IndexQueueClusterQueue).
    55  		WithIndex(&kueue.Workload{}, indexer.WorkloadQueueKey, indexer.IndexWorkloadQueue).
    56  		WithIndex(&kueue.Workload{}, indexer.WorkloadClusterQueueKey, indexer.IndexWorkloadClusterQueue)
    57  }
    58  
    59  type builderIndexer struct {
    60  	*fake.ClientBuilder
    61  }
    62  
    63  func (b *builderIndexer) IndexField(ctx context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error {
    64  	b.ClientBuilder = b.ClientBuilder.WithIndex(obj, field, extractValue)
    65  	return nil
    66  }
    67  
    68  func AsIndexer(builder *fake.ClientBuilder) client.FieldIndexer {
    69  	return &builderIndexer{ClientBuilder: builder}
    70  }
    71  
    72  type EventRecord struct {
    73  	Key       types.NamespacedName
    74  	EventType string
    75  	Reason    string
    76  	Message   string
    77  	// add annotations if ever needed
    78  }
    79  
    80  type EventRecorder struct {
    81  	lock           sync.Mutex
    82  	RecordedEvents []EventRecord
    83  }
    84  
    85  func SortEvents(ei, ej EventRecord) bool {
    86  	if cmp := strings.Compare(ei.Key.String(), ej.Key.String()); cmp != 0 {
    87  		return cmp < 0
    88  	}
    89  	if cmp := strings.Compare(ei.EventType, ej.EventType); cmp != 0 {
    90  		return cmp < 0
    91  	}
    92  	if cmp := strings.Compare(ei.Reason, ej.Reason); cmp != 0 {
    93  		return cmp < 0
    94  	}
    95  	if cmp := strings.Compare(ei.Message, ej.Message); cmp != 0 {
    96  		return cmp < 0
    97  	}
    98  	return false
    99  }
   100  
   101  func (tr *EventRecorder) Event(object runtime.Object, eventtype, reason, message string) {
   102  	tr.Eventf(object, eventtype, reason, message)
   103  }
   104  
   105  func (tr *EventRecorder) Eventf(object runtime.Object, eventtype, reason, messageFmt string, args ...interface{}) {
   106  	tr.AnnotatedEventf(object, nil, eventtype, reason, messageFmt, args...)
   107  }
   108  
   109  func (tr *EventRecorder) AnnotatedEventf(targetObject runtime.Object, annotations map[string]string, eventtype, reason, messageFmt string, args ...interface{}) {
   110  	tr.lock.Lock()
   111  	defer tr.lock.Unlock()
   112  	key := types.NamespacedName{}
   113  	if cobj, iscobj := targetObject.(client.Object); iscobj {
   114  		key = client.ObjectKeyFromObject(cobj)
   115  	}
   116  	tr.RecordedEvents = append(tr.RecordedEvents, EventRecord{
   117  		Key:       key,
   118  		EventType: eventtype,
   119  		Reason:    reason,
   120  		Message:   fmt.Sprintf(messageFmt, args...),
   121  	})
   122  }