sigs.k8s.io/kueue@v0.6.2/test/util/e2e.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  
     8  	"github.com/google/go-cmp/cmp/cmpopts"
     9  	"github.com/onsi/gomega"
    10  	appsv1 "k8s.io/api/apps/v1"
    11  	corev1 "k8s.io/api/core/v1"
    12  	"k8s.io/apimachinery/pkg/types"
    13  	"k8s.io/client-go/kubernetes/scheme"
    14  	"sigs.k8s.io/controller-runtime/pkg/client"
    15  	"sigs.k8s.io/controller-runtime/pkg/client/config"
    16  	jobset "sigs.k8s.io/jobset/api/jobset/v1alpha2"
    17  
    18  	kueuealpha "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
    19  	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
    20  	visibility "sigs.k8s.io/kueue/apis/visibility/v1alpha1"
    21  	kueueclientset "sigs.k8s.io/kueue/client-go/clientset/versioned"
    22  	visibilityv1alpha1 "sigs.k8s.io/kueue/client-go/clientset/versioned/typed/visibility/v1alpha1"
    23  )
    24  
    25  func CreateClientUsingCluster(kContext string) client.WithWatch {
    26  	cfg, err := config.GetConfigWithContext(kContext)
    27  	if err != nil {
    28  		fmt.Printf("unable to get kubeconfig for context %q: %s", kContext, err)
    29  		os.Exit(1)
    30  	}
    31  	gomega.ExpectWithOffset(1, cfg).NotTo(gomega.BeNil())
    32  
    33  	err = kueue.AddToScheme(scheme.Scheme)
    34  	gomega.ExpectWithOffset(1, err).NotTo(gomega.HaveOccurred())
    35  
    36  	err = kueuealpha.AddToScheme(scheme.Scheme)
    37  	gomega.ExpectWithOffset(1, err).NotTo(gomega.HaveOccurred())
    38  
    39  	err = visibility.AddToScheme(scheme.Scheme)
    40  	gomega.ExpectWithOffset(1, err).NotTo(gomega.HaveOccurred())
    41  
    42  	err = jobset.AddToScheme(scheme.Scheme)
    43  	gomega.ExpectWithOffset(1, err).NotTo(gomega.HaveOccurred())
    44  
    45  	client, err := client.NewWithWatch(cfg, client.Options{Scheme: scheme.Scheme})
    46  	gomega.ExpectWithOffset(1, err).NotTo(gomega.HaveOccurred())
    47  	return client
    48  }
    49  
    50  func CreateVisibilityClient(user string) visibilityv1alpha1.VisibilityV1alpha1Interface {
    51  	cfg, err := config.GetConfigWithContext("")
    52  	if err != nil {
    53  		fmt.Printf("unable to get kubeconfig: %s", err)
    54  		os.Exit(1)
    55  	}
    56  	gomega.ExpectWithOffset(1, cfg).NotTo(gomega.BeNil())
    57  
    58  	if user != "" {
    59  		cfg.Impersonate.UserName = user
    60  	}
    61  
    62  	kueueClient, err := kueueclientset.NewForConfig(cfg)
    63  	if err != nil {
    64  		fmt.Printf("unable to create kueue clientset: %s", err)
    65  		os.Exit(1)
    66  	}
    67  	visibilityClient := kueueClient.VisibilityV1alpha1()
    68  	return visibilityClient
    69  }
    70  
    71  func WaitForKueueAvailability(ctx context.Context, k8sClient client.Client) {
    72  	kcmKey := types.NamespacedName{
    73  		Namespace: "kueue-system",
    74  		Name:      "kueue-controller-manager",
    75  	}
    76  	deployment := &appsv1.Deployment{}
    77  	pods := corev1.PodList{}
    78  	gomega.EventuallyWithOffset(1, func(g gomega.Gomega) error {
    79  		g.Expect(k8sClient.Get(ctx, kcmKey, deployment)).To(gomega.Succeed())
    80  		g.Expect(k8sClient.List(ctx, &pods, client.InNamespace("kueue-system"), client.MatchingLabels(deployment.Spec.Selector.MatchLabels))).To(gomega.Succeed())
    81  		for _, pod := range pods.Items {
    82  			for _, cs := range pod.Status.ContainerStatuses {
    83  				if cs.RestartCount > 0 {
    84  					return gomega.StopTrying(fmt.Sprintf("%q in %q has restarted %d times", cs.Name, pod.Name, cs.RestartCount))
    85  				}
    86  			}
    87  		}
    88  		g.Expect(deployment.Status.Conditions).To(gomega.ContainElement(gomega.BeComparableTo(
    89  			appsv1.DeploymentCondition{
    90  				Type:   appsv1.DeploymentAvailable,
    91  				Status: corev1.ConditionTrue,
    92  			},
    93  			cmpopts.IgnoreFields(appsv1.DeploymentCondition{}, "Reason", "Message", "LastUpdateTime", "LastTransitionTime"))))
    94  		return nil
    95  
    96  	}, StartUpTimeout, Interval).Should(gomega.Succeed())
    97  }