From eaa0dc69b1243090d46ec39724b8ee95cc107f00 Mon Sep 17 00:00:00 2001
From: chatton <cianhatton@gmail.com>
Date: Wed, 24 Mar 2021 10:38:10 +0000
Subject: [PATCH] Added function for startup probe

---
 pkg/kube/container/container_test.go | 12 ++++++++++++
 pkg/kube/container/containers.go     | 14 ++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pkg/kube/container/container_test.go b/pkg/kube/container/container_test.go
index 04f9b000d..5c08b14b9 100644
--- a/pkg/kube/container/container_test.go
+++ b/pkg/kube/container/container_test.go
@@ -37,6 +37,13 @@ func TestContainer(t *testing.T) {
 			probes.WithFailureThreshold(15),
 			probes.WithPeriodSeconds(10),
 		)),
+		WithStartupProbe(
+			probes.Apply(
+				probes.WithExecCommand([]string{"startup-exec"}),
+				probes.WithFailureThreshold(20),
+				probes.WithPeriodSeconds(30),
+			),
+		),
 		WithResourceRequirements(resourcerequirements.Defaults()),
 		WithCommand([]string{"container-cmd"}),
 		WithEnvs(
@@ -71,6 +78,11 @@ func TestContainer(t *testing.T) {
 	assert.Equal(t, int32(10), liveNessProbe.PeriodSeconds)
 	assert.Equal(t, "liveness-exec", liveNessProbe.Exec.Command[0])
 
+	startupProbe := c.StartupProbe
+	assert.Equal(t, int32(20), startupProbe.FailureThreshold)
+	assert.Equal(t, int32(30), startupProbe.PeriodSeconds)
+	assert.Equal(t, "startup-exec", startupProbe.Exec.Command[0])
+
 	assert.Equal(t, c.Resources, resourcerequirements.Defaults())
 
 	assert.Len(t, c.Command, 1)
diff --git a/pkg/kube/container/containers.go b/pkg/kube/container/containers.go
index bd0b58890..94c5fc3b7 100644
--- a/pkg/kube/container/containers.go
+++ b/pkg/kube/container/containers.go
@@ -75,12 +75,22 @@ func WithReadinessProbe(probeFunc func(*corev1.Probe)) Modification {
 }
 
 // WithLivenessProbe modifies the container's Liveness Probe
-func WithLivenessProbe(readinessProbeFunc func(*corev1.Probe)) Modification {
+func WithLivenessProbe(livenessProbeFunc func(*corev1.Probe)) Modification {
 	return func(container *corev1.Container) {
 		if container.LivenessProbe == nil {
 			container.LivenessProbe = &corev1.Probe{}
 		}
-		readinessProbeFunc(container.LivenessProbe)
+		livenessProbeFunc(container.LivenessProbe)
+	}
+}
+
+// WithStartupProbe modifies the container's Startup Probe
+func WithStartupProbe(startupProbeFunc func(*corev1.Probe)) Modification {
+	return func(container *corev1.Container) {
+		if container.StartupProbe == nil {
+			container.StartupProbe = &corev1.Probe{}
+		}
+		startupProbeFunc(container.StartupProbe)
 	}
 }