feat(k8s): add WaitUntilDaemonSetAvailable helpers (#1815)

* feat(k8s): add WaitUntilDaemonSetAvailable helpers

Mirror the deployment-wait pattern for DaemonSet: Context/ContextE plus
deprecated non-context variants, IsDaemonSetAvailable for status checks,
and a DaemonSetNotAvailable error. Availability is gated on observed
generation matching spec generation, DesiredNumberScheduled > 0, and
NumberAvailable == DesiredNumberScheduled, since DaemonSetCondition is
not always populated.

* review: align IsDaemonSetAvailable with kubectl rollout status

- Drop the DesiredNumberScheduled == 0 gate. A DaemonSet whose selector
  matches no nodes is "rolled out" per kubectl; the gate caused a wait
  to spin until the retry budget was exhausted in that case.
- Add UpdatedNumberScheduled == DesiredNumberScheduled to the readiness
  predicate, mirroring kubectl's daemonset rollout-status logic.
  Otherwise the wait could return success during a rolling update,
  while old pods still satisfy availability.
- Expand DaemonSetNotAvailable.Error to include UpdatedNumberScheduled,
  NumberReady, and NumberMisscheduled so timeout messages identify
  whether the rollout is stuck on update progress, readiness, or
  misscheduled pods.
- Add control-plane toleration alongside the existing master toleration
  so the integration test schedules on Kubernetes 1.24+ clusters.
- Cover the new behavior with unit-test cases for zero-desired,
  mid-rolling-update, and stale-observed-generation states.

* lint: satisfy wsl_v5 and govet fieldalignment

- Add the required blank line between the back-to-back if statements
  in IsDaemonSetAvailable.
- Inline err.daemonSet.Status references in DaemonSetNotAvailable.Error
  rather than aliasing to a local var, since wsl_v5 flags the cuddled
  intermediate assign+return.
- Reorder the IsDaemonSetAvailable test case struct so the pointer
  field precedes string and bool, satisfying govet fieldalignment.

Author: james00012

modules/k8s/daemonset.go

@@ -2,11 +2,14 @@ package k8s //nolint:dupl // structural pattern for k8s resource operations
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	appsv1 "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	"github.com/gruntwork-io/terratest/modules/retry"
 	"github.com/gruntwork-io/terratest/modules/testing"
 )
 
@@ -100,3 +103,99 @@ func GetDaemonSet(t testing.TestingT, options *KubectlOptions, daemonSetName str
 func GetDaemonSetE(t testing.TestingT, options *KubectlOptions, daemonSetName string) (*appsv1.DaemonSet, error) {
 	return GetDaemonSetContextE(t, context.Background(), options, daemonSetName)
 }
+
+// WaitUntilDaemonSetAvailableContextE waits until all desired pods of the daemonset are available on their nodes,
+// retrying the check for the specified amount of times, sleeping for the provided duration between each try.
+// The ctx parameter supports cancellation and timeouts.
+func WaitUntilDaemonSetAvailableContextE( //nolint:dupl // similar retry pattern across resource types is intentional
+	t testing.TestingT,
+	ctx context.Context,
+	options *KubectlOptions,
+	daemonSetName string,
+	retries int,
+	sleepBetweenRetries time.Duration,
+) error {
+	statusMsg := fmt.Sprintf("Wait for daemonset %s to be provisioned.", daemonSetName)
+
+	message, err := retry.DoWithRetryContextE(
+		t,
+		ctx,
+		statusMsg,
+		retries,
+		sleepBetweenRetries,
+		func() (string, error) {
+			daemonSet, err := GetDaemonSetContextE(t, ctx, options, daemonSetName)
+			if err != nil {
+				return "", err
+			}
+
+			if !IsDaemonSetAvailable(daemonSet) {
+				return "", NewDaemonSetNotAvailableError(daemonSet)
+			}
+
+			return "DaemonSet is now available", nil
+		},
+	)
+	if err != nil {
+		options.Logger.Logf(t, "Timedout waiting for DaemonSet to be provisioned: %s", err)
+		return err
+	}
+
+	options.Logger.Logf(t, "%s", message)
+
+	return nil
+}
+
+// WaitUntilDaemonSetAvailableContext waits until all desired pods of the daemonset are available on their nodes,
+// retrying the check for the specified amount of times, sleeping for the provided duration between each try.
+// The ctx parameter supports cancellation and timeouts.
+// This will fail the test if there is an error.
+func WaitUntilDaemonSetAvailableContext(t testing.TestingT, ctx context.Context, options *KubectlOptions, daemonSetName string, retries int, sleepBetweenRetries time.Duration) {
+	t.Helper()
+	require.NoError(t, WaitUntilDaemonSetAvailableContextE(t, ctx, options, daemonSetName, retries, sleepBetweenRetries))
+}
+
+// WaitUntilDaemonSetAvailable waits until all desired pods of the daemonset are available on their nodes,
+// retrying the check for the specified amount of times, sleeping
+// for the provided duration between each try.
+// This will fail the test if there is an error.
+//
+// Deprecated: Use [WaitUntilDaemonSetAvailableContext] instead.
+func WaitUntilDaemonSetAvailable(t testing.TestingT, options *KubectlOptions, daemonSetName string, retries int, sleepBetweenRetries time.Duration) {
+	t.Helper()
+	WaitUntilDaemonSetAvailableContext(t, context.Background(), options, daemonSetName, retries, sleepBetweenRetries)
+}
+
+// WaitUntilDaemonSetAvailableE waits until all desired pods of the daemonset are available on their nodes,
+// retrying the check for the specified amount of times, sleeping
+// for the provided duration between each try.
+//
+// Deprecated: Use [WaitUntilDaemonSetAvailableContextE] instead.
+func WaitUntilDaemonSetAvailableE(
+	t testing.TestingT,
+	options *KubectlOptions,
+	daemonSetName string,
+	retries int,
+	sleepBetweenRetries time.Duration,
+) error {
+	return WaitUntilDaemonSetAvailableContextE(t, context.Background(), options, daemonSetName, retries, sleepBetweenRetries)
+}
+
+// IsDaemonSetAvailable returns true once the daemonset's rollout is complete. The check mirrors `kubectl rollout
+// status ds`: the controller has observed the latest spec, every scheduled pod has been updated to the current
+// generation, and every desired pod is available. Status fields are used directly rather than DaemonSetCondition
+// because the controller does not always populate that field.
+//
+// A daemonset whose node selector matches zero nodes (DesiredNumberScheduled == 0) is treated as available — this
+// matches the kubectl behavior where such a daemonset is considered "successfully rolled out".
+func IsDaemonSetAvailable(ds *appsv1.DaemonSet) bool {
+	if ds.Status.ObservedGeneration < ds.Generation {
+		return false
+	}
+
+	if ds.Status.UpdatedNumberScheduled < ds.Status.DesiredNumberScheduled {
+		return false
+	}
+
+	return ds.Status.NumberAvailable >= ds.Status.DesiredNumberScheduled
+}

modules/k8s/daemonset_test.go

@@ -1,5 +1,5 @@
-//go:build kubernetes
-// +build kubernetes
+//go:build kubeall || kubernetes
+// +build kubeall kubernetes
 
 // NOTE: we have build tags to differentiate kubernetes tests from non-kubernetes tests. This is done because minikube
 // is heavy and can interfere with docker related tests in terratest. Specifically, many of the tests start to fail with
@@ -10,9 +10,11 @@ package k8s_test
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/gruntwork-io/terratest/modules/k8s"
 
+	appsv1 "k8s.io/api/apps/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"strings"
@@ -63,6 +65,105 @@ func TestListDaemonSetsReturnsCorrectServiceInCorrectNamespace(t *testing.T) {
 	require.Equal(t, daemonSet.Namespace, uniqueID)
 }
 
+func TestWaitUntilDaemonSetAvailable(t *testing.T) {
+	t.Parallel()
+
+	uniqueID := strings.ToLower(random.UniqueID())
+	options := k8s.NewKubectlOptions("", "", uniqueID)
+	configData := fmt.Sprintf(exampleDaemonSetYAMLTemplate, uniqueID, uniqueID)
+
+	k8s.KubectlApplyFromString(t, options, configData)
+	defer k8s.KubectlDeleteFromString(t, options, configData)
+
+	k8s.WaitUntilDaemonSetAvailable(t, options, "sample-ds", 60, 1*time.Second)
+}
+
+func TestIsDaemonSetAvailable(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		ds             *appsv1.DaemonSet
+		title          string
+		expectedResult bool
+	}{
+		{
+			title: "AvailableWhenAllPodsUpdatedAndAvailable",
+			ds: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 1},
+				Status: appsv1.DaemonSetStatus{
+					ObservedGeneration:     1,
+					DesiredNumberScheduled: 3,
+					UpdatedNumberScheduled: 3,
+					NumberAvailable:        3,
+				},
+			},
+			expectedResult: true,
+		},
+		{
+			title: "AvailableWhenNoNodesMatchSelector",
+			ds: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 1},
+				Status: appsv1.DaemonSetStatus{
+					ObservedGeneration:     1,
+					DesiredNumberScheduled: 0,
+					UpdatedNumberScheduled: 0,
+					NumberAvailable:        0,
+				},
+			},
+			expectedResult: true,
+		},
+		{
+			title: "NotAvailableWhenSomePodsNotYetAvailable",
+			ds: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 1},
+				Status: appsv1.DaemonSetStatus{
+					ObservedGeneration:     1,
+					DesiredNumberScheduled: 3,
+					UpdatedNumberScheduled: 3,
+					NumberAvailable:        2,
+				},
+			},
+			expectedResult: false,
+		},
+		{
+			title: "NotAvailableMidRollingUpdate",
+			ds: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 2},
+				Status: appsv1.DaemonSetStatus{
+					ObservedGeneration:     2,
+					DesiredNumberScheduled: 3,
+					UpdatedNumberScheduled: 1,
+					NumberAvailable:        3,
+				},
+			},
+			expectedResult: false,
+		},
+		{
+			title: "NotAvailableWhenObservedGenerationStale",
+			ds: &appsv1.DaemonSet{
+				ObjectMeta: metav1.ObjectMeta{Generation: 2},
+				Status: appsv1.DaemonSetStatus{
+					ObservedGeneration:     1,
+					DesiredNumberScheduled: 3,
+					UpdatedNumberScheduled: 3,
+					NumberAvailable:        3,
+				},
+			},
+			expectedResult: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.title, func(t *testing.T) {
+			t.Parallel()
+
+			actualResult := k8s.IsDaemonSetAvailable(tc.ds)
+			require.Equal(t, tc.expectedResult, actualResult)
+		})
+	}
+}
+
 const exampleDaemonSetYAMLTemplate = `---
 apiVersion: v1
 kind: Namespace
@@ -88,6 +189,8 @@ spec:
       tolerations:
       - key: node-role.kubernetes.io/master
         effect: NoSchedule
+      - key: node-role.kubernetes.io/control-plane
+        effect: NoSchedule
       containers:
       - name: alpine
         image: alpine:3.8

modules/k8s/errors.go

@@ -111,6 +111,33 @@ func NewDeploymentNotAvailableError(deploy *appsv1.Deployment) DeploymentNotAvai
 	return DeploymentNotAvailable{deploy}
 }
 
+// DaemonSetNotAvailable is returned when a Kubernetes daemonset has not yet rolled out the desired number of pods.
+type DaemonSetNotAvailable struct {
+	daemonSet *appsv1.DaemonSet
+}
+
+// Error is a simple function to return a formatted error message as a string
+func (err DaemonSetNotAvailable) Error() string {
+	return fmt.Sprintf(
+		"DaemonSet %s is not available: generation observed %d/%d, updated %d/%d, ready %d/%d, available %d/%d, misscheduled %d",
+		err.daemonSet.Name,
+		err.daemonSet.Status.ObservedGeneration,
+		err.daemonSet.Generation,
+		err.daemonSet.Status.UpdatedNumberScheduled,
+		err.daemonSet.Status.DesiredNumberScheduled,
+		err.daemonSet.Status.NumberReady,
+		err.daemonSet.Status.DesiredNumberScheduled,
+		err.daemonSet.Status.NumberAvailable,
+		err.daemonSet.Status.DesiredNumberScheduled,
+		err.daemonSet.Status.NumberMisscheduled,
+	)
+}
+
+// NewDaemonSetNotAvailableError returns a DaemonSetNotAvailable struct when Kubernetes deems a daemonset is not available
+func NewDaemonSetNotAvailableError(ds *appsv1.DaemonSet) DaemonSetNotAvailable {
+	return DaemonSetNotAvailable{ds}
+}
+
 // PodNotAvailable is returned when a Kubernetes service is not yet available to accept traffic.
 type PodNotAvailable struct {
 	pod *corev1.Pod