TOFIX: reduce diff

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

tests/e2e/device-plugin_test.go

@@ -32,7 +32,6 @@ import (
 	"k8s.io/apimachinery/pkg/util/rand"
 
 	"github.com/NVIDIA/k8s-device-plugin/tests/e2e/common/diagnostics"
-	"github.com/NVIDIA/k8s-device-plugin/tests/e2e/internal"
 )
 
 const (
@@ -96,7 +95,7 @@ var _ = Describe("GPU Device Plugin", Ordered, Label("gpu", "e2e", "device-plugi
 		// Note: DaemonSet names are dynamically generated with the Helm release prefix,
 		// so we wait for all DaemonSets in the namespace rather than specific names
 		By("Waiting for all DaemonSets to be ready")
-		err = internal.WaitForDaemonSetsReady(ctx, clientSet, testNamespace.Name, "app.kubernetes.io/name=nvidia-device-plugin")
+		err = waitForDaemonSetsReady(ctx, clientSet, testNamespace.Name, "app.kubernetes.io/name=nvidia-device-plugin")
 		Expect(err).NotTo(HaveOccurred())
 	})
 

tests/e2e/e2e_test.go

@@ -874,3 +874,45 @@ func getEnvVarOrDefault[T any](key string, defaultValue T) T {
 	}
 	return val
 }
+
+// waitForDaemonSetsReady waits for DaemonSets in a namespace to be ready, optionally filtered by label selector
+func waitForDaemonSetsReady(ctx context.Context, client kubernetes.Interface, namespace, labelSelector string) error {
+	EventuallyWithOffset(1, func(g Gomega) error {
+		dsList, err := client.AppsV1().DaemonSets(namespace).List(ctx, metav1.ListOptions{
+			LabelSelector: labelSelector,
+		})
+		if err != nil {
+			return err
+		}
+
+		if len(dsList.Items) == 0 {
+			return fmt.Errorf("no daemonsets found in namespace %s with selector '%s'", namespace, labelSelector)
+		}
+
+		for _, ds := range dsList.Items {
+			// Skip if no pods are desired
+			if ds.Status.DesiredNumberScheduled == 0 {
+				continue
+			}
+
+			if ds.Status.NumberReady != ds.Status.DesiredNumberScheduled {
+				return fmt.Errorf("daemonset %s/%s rollout incomplete: %d/%d pods ready",
+					namespace, ds.Name, ds.Status.NumberReady, ds.Status.DesiredNumberScheduled)
+			}
+
+			if ds.Status.UpdatedNumberScheduled != ds.Status.DesiredNumberScheduled {
+				return fmt.Errorf("daemonset %s/%s update incomplete: %d/%d pods updated",
+					namespace, ds.Name, ds.Status.UpdatedNumberScheduled, ds.Status.DesiredNumberScheduled)
+			}
+
+			// Check generation to ensure we're looking at the latest spec
+			if ds.Generation != ds.Status.ObservedGeneration {
+				return fmt.Errorf("daemonset %s/%s generation mismatch: %d != %d",
+					namespace, ds.Name, ds.Generation, ds.Status.ObservedGeneration)
+			}
+		}
+
+		return nil
+	}).WithContext(ctx).WithPolling(2 * time.Second).WithTimeout(5 * time.Minute).Should(Succeed())
+	return nil
+}

tests/e2e/gpu-feature-discovery_test.go

@@ -32,7 +32,6 @@ import (
 	"k8s.io/apimachinery/pkg/util/rand"
 
 	"github.com/NVIDIA/k8s-device-plugin/tests/e2e/common/diagnostics"
-	"github.com/NVIDIA/k8s-device-plugin/tests/e2e/internal"
 )
 
 var expectedLabelPatterns = k8sLabels{
@@ -117,7 +116,7 @@ var _ = Describe("GPU Feature Discovery", Ordered, Label("gfd", "gpu", "e2e"), f
 		// Note: DaemonSet names are dynamically generated with the Helm release prefix,
 		// so we wait for all DaemonSets in the namespace rather than specific names
 		By("Waiting for all DaemonSets to be ready")
-		err = internal.WaitForDaemonSetsReady(ctx, clientSet, testNamespace.Name, "app.kubernetes.io/name=nvidia-device-plugin")
+		err = waitForDaemonSetsReady(ctx, clientSet, testNamespace.Name, "app.kubernetes.io/name=nvidia-device-plugin")
 		Expect(err).NotTo(HaveOccurred())
 	})