Refactor e2e test infrastructure with unified helpers leveraging GPU operator clientsets

Signed-off-by: Karthik Vetrivel <[email protected]>

Author: karthikvetrivel

tests/e2e/gpu_operator_test.go

@@ -28,8 +28,7 @@ import (
 
 	"github.com/NVIDIA/gpu-operator/tests/e2e/framework"
 	e2elog "github.com/NVIDIA/gpu-operator/tests/e2e/framework/logs"
-	k8stest "github.com/NVIDIA/gpu-operator/tests/e2e/kubernetes"
-	"github.com/NVIDIA/gpu-operator/tests/e2e/operator"
+	"github.com/NVIDIA/gpu-operator/tests/e2e/helpers"
 )
 
 var _ = Describe(e2eTestPrefix+"-premerge-suite", func() {
@@ -44,15 +43,15 @@ var _ = Describe(e2eTestPrefix+"-premerge-suite", func() {
 
 			// Init global suite vars vars
 			var (
-				operatorClient  *operator.Client
+				operatorClient  *helpers.OperatorClient
 				helmReleaseName string
-				k8sClient       *k8stest.Client
+				k8sClient       *helpers.PodClient
 				testNamespace   *corev1.Namespace
 			)
 
 			BeforeAll(func(ctx context.Context) {
 				var err error
-				k8sClient = k8stest.NewClient(f.ClientSet.CoreV1())
+				k8sClient = helpers.NewPodClient(f.ClientSet.CoreV1())
 				nsLabels := map[string]string{
 					"e2e-run": string(framework.RunID),
 				}
@@ -62,10 +61,10 @@ var _ = Describe(e2eTestPrefix+"-premerge-suite", func() {
 					Fail(fmt.Sprintf("failed to create gpu operator namespace %s: %v", tcfg.namespace, err))
 				}
 
-				operatorClient, err = operator.NewClient(
-					operator.WithNamespace(testNamespace.Name),
-					operator.WithKubeConfig(framework.TestContext.KubeConfig),
-					operator.WithChart(tcfg.helmChart),
+				operatorClient, err = helpers.NewOperatorClient(
+					helpers.WithNamespace(testNamespace.Name),
+					helpers.WithKubeConfig(framework.TestContext.KubeConfig),
+					helpers.WithChart(tcfg.helmChart),
 				)
 				if err != nil {
 					Fail(fmt.Sprintf("failed to instantiate gpu operator client: %v", err))
@@ -79,7 +78,7 @@ var _ = Describe(e2eTestPrefix+"-premerge-suite", func() {
 					fmt.Sprintf("validator.image=%s", tcfg.validatorImage),
 					fmt.Sprintf("validator.version=%s", tcfg.validatorVersion),
 				}
-				helmReleaseName, err = operatorClient.Install(ctx, values, operator.ChartOptions{
+				helmReleaseName, err = operatorClient.Install(ctx, values, helpers.ChartOptions{
 					CleanupOnFail: true,
 					GenerateName:  true,
 					Timeout:       5 * time.Minute,
@@ -157,8 +156,12 @@ var _ = Describe(e2eTestPrefix+"-premerge-suite", func() {
 						hasRestarts, err := k8sClient.EnsureNoPodRestarts(ctx, pod.Name, pod.Namespace)
 						Expect(err).NotTo(HaveOccurred())
 						if !hasRestarts {
-							errLogs := k8sClient.GetPodLogs(ctx, pod)
-							e2elog.Logf("printing logs from the pod %s/%s: %s", pod.Namespace, pod.Name, errLogs)
+							errLogs, err := k8sClient.GetPodLogs(ctx, pod)
+							if err != nil {
+								e2elog.Logf("WARN: failed to retrieve logs from pod %s/%s: %v", pod.Namespace, pod.Name, err)
+							} else {
+								e2elog.Logf("printing logs from the pod %s/%s: %s", pod.Namespace, pod.Name, errLogs)
+							}
 							e2elog.Failf("pod %s/%s has unexpected restarts", pod.Namespace, pod.Name)
 						}
 					}

tests/e2e/helpers/clusterpolicy.go

@@ -0,0 +1,130 @@
+/**
+# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+#
+# 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 helpers
+
+import (
+	"context"
+	"time"
+
+	nvidiav1 "github.com/NVIDIA/gpu-operator/api/nvidia/v1"
+	gpuclientset "github.com/NVIDIA/gpu-operator/api/versioned"
+	"github.com/NVIDIA/gpu-operator/internal/conditions"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/util/retry"
+	"k8s.io/utils/ptr"
+)
+
+type ClusterPolicyHelper struct {
+	client gpuclientset.Interface
+}
+
+func NewClusterPolicyHelper(client gpuclientset.Interface) *ClusterPolicyHelper {
+	return &ClusterPolicyHelper{
+		client: client,
+	}
+}
+
+func (h *ClusterPolicyHelper) Get(ctx context.Context, name string) (*nvidiav1.ClusterPolicy, error) {
+	return h.client.NvidiaV1().ClusterPolicies().Get(ctx, name, metav1.GetOptions{})
+}
+
+func (h *ClusterPolicyHelper) Update(ctx context.Context, cp *nvidiav1.ClusterPolicy) (*nvidiav1.ClusterPolicy, error) {
+	return h.client.NvidiaV1().ClusterPolicies().Update(ctx, cp, metav1.UpdateOptions{})
+}
+
+// modify applies a mutation function to a ClusterPolicy and persists the changes.
+// It uses RetryOnConflict to handle concurrent modifications by the operator controller.
+func (h *ClusterPolicyHelper) modify(ctx context.Context, name string, mutate func(*nvidiav1.ClusterPolicy)) error {
+	return retry.RetryOnConflict(retry.DefaultBackoff, func() error {
+		clusterPolicy, err := h.Get(ctx, name)
+		if err != nil {
+			return err
+		}
+
+		mutate(clusterPolicy)
+
+		_, err = h.Update(ctx, clusterPolicy)
+		return err
+	})
+}
+
+func (h *ClusterPolicyHelper) UpdateDriverVersion(ctx context.Context, name, version string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.Driver.Version = version
+	})
+}
+
+func (h *ClusterPolicyHelper) EnableDCGM(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.DCGM.Enabled = ptr.To(true)
+	})
+}
+
+func (h *ClusterPolicyHelper) DisableDCGM(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.DCGM.Enabled = ptr.To(false)
+	})
+}
+
+func (h *ClusterPolicyHelper) EnableDCGMExporter(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.DCGMExporter.Enabled = ptr.To(true)
+	})
+}
+
+func (h *ClusterPolicyHelper) DisableDCGMExporter(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.DCGMExporter.Enabled = ptr.To(false)
+	})
+}
+
+func (h *ClusterPolicyHelper) EnableGFD(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.GPUFeatureDiscovery.Enabled = ptr.To(true)
+	})
+}
+
+func (h *ClusterPolicyHelper) DisableGFD(ctx context.Context, name string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.GPUFeatureDiscovery.Enabled = ptr.To(false)
+	})
+}
+
+func (h *ClusterPolicyHelper) SetMIGStrategy(ctx context.Context, name, strategy string) error {
+	return h.modify(ctx, name, func(clusterPolicy *nvidiav1.ClusterPolicy) {
+		clusterPolicy.Spec.MIG.Strategy = nvidiav1.MIGStrategy(strategy)
+	})
+}
+
+func (h *ClusterPolicyHelper) WaitForReady(ctx context.Context, name string, timeout time.Duration) error {
+	return wait.PollUntilContextTimeout(ctx, defaultPollingInterval, timeout, true, func(ctx context.Context) (bool, error) {
+		clusterPolicy, err := h.Get(ctx, name)
+		if err != nil {
+			return false, err
+		}
+
+		for _, condition := range clusterPolicy.Status.Conditions {
+			if condition.Type == conditions.Ready && condition.Status == metav1.ConditionTrue {
+				return true, nil
+			}
+		}
+
+		return false, nil
+	})
+}
+

tests/e2e/helpers/constants.go

@@ -0,0 +1,28 @@
+/**
+# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+#
+# 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 helpers
+
+import "time"
+
+const (
+	// defaultPollingInterval is the default interval for polling operations
+	defaultPollingInterval = 5 * time.Second
+
+	// upgradeDoneState represents the state when a driver upgrade is complete
+	upgradeDoneState = "upgrade-done"
+)
+

tests/e2e/helpers/daemonset.go

@@ -0,0 +1,131 @@
+/**
+# Copyright (c) NVIDIA CORPORATION.  All rights reserved.
+#
+# 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 helpers
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	appsv1 "k8s.io/api/apps/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/kubernetes"
+)
+
+type DaemonSetHelper struct {
+	client kubernetes.Interface
+}
+
+func NewDaemonSetHelper(client kubernetes.Interface) *DaemonSetHelper {
+	return &DaemonSetHelper{
+		client: client,
+	}
+}
+
+func (h *DaemonSetHelper) GetByLabel(ctx context.Context, namespace, labelKey, labelValue string) (*appsv1.DaemonSet, error) {
+	labelSelector := labels.SelectorFromSet(map[string]string{
+		labelKey: labelValue,
+	}).String()
+
+	daemonSetList, err := h.client.AppsV1().DaemonSets(namespace).List(ctx, metav1.ListOptions{
+		LabelSelector: labelSelector,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to list DaemonSets: %w", err)
+	}
+
+	if len(daemonSetList.Items) == 0 {
+		return nil, fmt.Errorf("no DaemonSet found with label %s=%s", labelKey, labelValue)
+	}
+
+	if len(daemonSetList.Items) > 1 {
+		return nil, fmt.Errorf("multiple DaemonSets found with label %s=%s", labelKey, labelValue)
+	}
+
+	return &daemonSetList.Items[0], nil
+}
+
+func (h *DaemonSetHelper) Get(ctx context.Context, namespace, name string) (*appsv1.DaemonSet, error) {
+	return h.client.AppsV1().DaemonSets(namespace).Get(ctx, name, metav1.GetOptions{})
+}
+
+func (h *DaemonSetHelper) WaitForReady(ctx context.Context, namespace, name string, timeout time.Duration) error {
+	return wait.PollUntilContextTimeout(ctx, defaultPollingInterval, timeout, true, func(ctx context.Context) (bool, error) {
+		daemonSet, err := h.Get(ctx, namespace, name)
+		if err != nil {
+			return false, err
+		}
+
+		if daemonSet.Status.NumberReady == daemonSet.Status.DesiredNumberScheduled &&
+			daemonSet.Status.NumberReady > 0 {
+			return true, nil
+		}
+
+		return false, nil
+	})
+}
+
+func (h *DaemonSetHelper) IsReady(ctx context.Context, namespace, name string) (bool, error) {
+	daemonSet, err := h.Get(ctx, namespace, name)
+	if err != nil {
+		return false, err
+	}
+
+	return daemonSet.Status.NumberReady == daemonSet.Status.DesiredNumberScheduled && daemonSet.Status.NumberReady > 0, nil
+}
+
+func (h *DaemonSetHelper) GetImage(ctx context.Context, namespace, name string) (string, error) {
+	daemonSet, err := h.Get(ctx, namespace, name)
+	if err != nil {
+		return "", fmt.Errorf("failed to get DaemonSet: %w", err)
+	}
+
+	if len(daemonSet.Spec.Template.Spec.Containers) == 0 {
+		return "", fmt.Errorf("DaemonSet has no containers")
+	}
+
+	return daemonSet.Spec.Template.Spec.Containers[0].Image, nil
+}
+
+func (h *DaemonSetHelper) CheckNoRestarts(ctx context.Context, namespace, name string) error {
+	daemonSet, err := h.Get(ctx, namespace, name)
+	if err != nil {
+		return fmt.Errorf("failed to get DaemonSet: %w", err)
+	}
+
+	labelSelector := labels.SelectorFromSet(daemonSet.Spec.Selector.MatchLabels).String()
+	podList, err := h.client.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
+		LabelSelector: labelSelector,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to list pods: %w", err)
+	}
+
+	for _, pod := range podList.Items {
+		for _, containerStatus := range pod.Status.ContainerStatuses {
+			if containerStatus.RestartCount > 0 {
+				return fmt.Errorf("pod %s/%s container %s has %d restarts",
+					pod.Namespace, pod.Name, containerStatus.Name, containerStatus.RestartCount)
+			}
+		}
+	}
+
+	return nil
+}
+