Add liveness and startup probe helpers (elastic#8634)

Extend the PodTemplateBuilder with the methods WithLivenessProbe and WithStartupProbe.

Author: afharo

pkg/controller/common/container/defaulter.go

@@ -40,7 +40,9 @@ func (d Defaulter) From(other corev1.Container) Defaulter {
 		WithEnv(other.Env).
 		WithResources(other.Resources).
 		WithVolumeMounts(other.VolumeMounts).
-		WithReadinessProbe(other.ReadinessProbe)
+		WithReadinessProbe(other.ReadinessProbe).
+		WithLivenessProbe(other.LivenessProbe).
+		WithStartupProbe(other.StartupProbe)
 }
 
 func (d Defaulter) WithCommand(command []string) Defaulter {
@@ -96,6 +98,20 @@ func (d Defaulter) WithReadinessProbe(readinessProbe *corev1.Probe) Defaulter {
 	return d
 }
 
+func (d Defaulter) WithLivenessProbe(livenessProbe *corev1.Probe) Defaulter {
+	if d.base.LivenessProbe == nil {
+		d.base.LivenessProbe = livenessProbe
+	}
+	return d
+}
+
+func (d Defaulter) WithStartupProbe(startupProbe *corev1.Probe) Defaulter {
+	if d.base.StartupProbe == nil {
+		d.base.StartupProbe = startupProbe
+	}
+	return d
+}
+
 // envExists checks if an env var with the given name already exists in the provided slice.
 func (d Defaulter) envExists(name string) bool {
 	for _, v := range d.base.Env {

pkg/controller/common/defaults/pod_template.go

@@ -121,6 +121,18 @@ func (b *PodTemplateBuilder) WithReadinessProbe(readinessProbe corev1.Probe) *Po
 	return b
 }
 
+// WithLivenessProbe sets up the given liveness probe, unless already provided in the template.
+func (b *PodTemplateBuilder) WithLivenessProbe(livenessProbe corev1.Probe) *PodTemplateBuilder {
+	b.containerDefaulter.WithLivenessProbe(&livenessProbe)
+	return b
+}
+
+// WithStartupProbe sets up the given startup probe, unless already provided in the template.
+func (b *PodTemplateBuilder) WithStartupProbe(startupProbe corev1.Probe) *PodTemplateBuilder {
+	b.containerDefaulter.WithStartupProbe(&startupProbe)
+	return b
+}
+
 // WithAffinity sets a default affinity, unless already provided in the template.
 // An empty affinity in the spec is not overridden.
 func (b *PodTemplateBuilder) WithAffinity(affinity *corev1.Affinity) *PodTemplateBuilder {

pkg/controller/common/defaults/pod_template_test.go

@@ -274,6 +274,146 @@ func TestPodTemplateBuilder_WithReadinessProbe(t *testing.T) {
 	}
 }
 
+func TestPodTemplateBuilder_WithLivenessProbe(t *testing.T) {
+	containerName := "mycontainer"
+	tests := []struct {
+		name          string
+		PodTemplate   corev1.PodTemplateSpec
+		livenessProbe corev1.Probe
+		want          *corev1.Probe
+	}{
+		{
+			name:        "no liveness probe in pod template: use default one",
+			PodTemplate: corev1.PodTemplateSpec{},
+			livenessProbe: corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+			want: &corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+		},
+		{
+			name: "don't override pod template liveness probe",
+			PodTemplate: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: containerName,
+							LivenessProbe: &corev1.Probe{
+								ProbeHandler: corev1.ProbeHandler{
+									HTTPGet: &corev1.HTTPGetAction{
+										Path: "/user-provided",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			livenessProbe: corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+			want: &corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/user-provided",
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := NewPodTemplateBuilder(tt.PodTemplate, containerName)
+			if got := b.WithLivenessProbe(tt.livenessProbe).containerDefaulter.Container().LivenessProbe; !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("PodTemplateBuilder.WithLivenessProbe() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestPodTemplateBuilder_WithStartupProbe(t *testing.T) {
+	containerName := "mycontainer"
+	tests := []struct {
+		name         string
+		PodTemplate  corev1.PodTemplateSpec
+		startupProbe corev1.Probe
+		want         *corev1.Probe
+	}{
+		{
+			name:        "no startup probe in pod template: use default one",
+			PodTemplate: corev1.PodTemplateSpec{},
+			startupProbe: corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+			want: &corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+		},
+		{
+			name: "don't override pod template startup probe",
+			PodTemplate: corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: containerName,
+							StartupProbe: &corev1.Probe{
+								ProbeHandler: corev1.ProbeHandler{
+									HTTPGet: &corev1.HTTPGetAction{
+										Path: "/user-provided",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			startupProbe: corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/probe",
+					},
+				},
+			},
+			want: &corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{
+						Path: "/user-provided",
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := NewPodTemplateBuilder(tt.PodTemplate, containerName)
+			if got := b.WithStartupProbe(tt.startupProbe).containerDefaulter.Container().StartupProbe; !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("PodTemplateBuilder.WithStartupProbe() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func TestPodTemplateBuilder_WithAffinity(t *testing.T) {
 	defaultAffinity := &corev1.Affinity{
 		NodeAffinity: &corev1.NodeAffinity{},