Allow overriding readinessProbe and LivenessProbe (#1776)

There can only be 1 probe and without this change,
specifying an exec readinessProbe in the spec.override would
add it on top of the default TCP probe.

As for livenessProbe, since there's currently no default,
everything works without this change. However, we could add
the default livenessProbe in the future and allowing this override
doesn't break anything right now, so it's just future-proofing.

Co-authored-by: joey <[email protected]>

Author: mkuratczyk

internal/resource/statefulset.go

@@ -286,6 +286,15 @@ func patchPodSpec(podSpec, podSpecOverride *corev1.PodSpec) (corev1.PodSpec, err
 		sortVolumeMounts(patchedPodSpec.Containers[0].VolumeMounts)
 	}
 
+	// allow overriding the liveness and readiness probes
+	// note: as of 2024, we don't set a default LivenessProbe
+	if rmqContainer.LivenessProbe != nil {
+		patchedPodSpec.Containers[0].LivenessProbe = rmqContainer.LivenessProbe
+	}
+	if rmqContainer.ReadinessProbe != nil {
+		patchedPodSpec.Containers[0].ReadinessProbe = rmqContainer.ReadinessProbe
+	}
+
 	// A user may wish to override the controller-set securityContext for the RabbitMQ & init containers so that the
 	// container runtime can override them. If the securityContext has been set to an empty struct, `strategicpatch.StrategicMergePatch`
 	// won't pick this up, so manually override it here.

internal/resource/statefulset_test.go

@@ -2126,6 +2126,84 @@ default_pass = {{ .Data.data.password }}
 
 				})
 
+				It("can replace the default readinessProbe", func() {
+					instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{
+						Spec: &rabbitmqv1beta1.StatefulSetSpec{
+							Template: &rabbitmqv1beta1.PodTemplateSpec{
+								Spec: &corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name: "rabbitmq",
+											ReadinessProbe: &corev1.Probe{
+												ProbeHandler: corev1.ProbeHandler{
+													Exec: &corev1.ExecAction{
+														Command: []string{"custom-readiness-probe", "arg1"},
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					}
+
+					builder = &resource.RabbitmqResourceBuilder{
+						Instance: &instance,
+						Scheme:   scheme,
+					}
+					stsBuilder := builder.StatefulSet()
+					Expect(stsBuilder.Update(statefulSet)).To(Succeed())
+
+					Expect(statefulSet.Spec.Template.Spec.Containers[0].ReadinessProbe.ProbeHandler).To(Equal(
+						corev1.ProbeHandler{
+							Exec: &corev1.ExecAction{
+								Command: []string{"custom-readiness-probe", "arg1"},
+							},
+						},
+					))
+				})
+
+				It("can add/replace a LivenessProbe", func() {
+					// note: we currently don't have a default LivenessProbe
+					instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{
+						Spec: &rabbitmqv1beta1.StatefulSetSpec{
+							Template: &rabbitmqv1beta1.PodTemplateSpec{
+								Spec: &corev1.PodSpec{
+									Containers: []corev1.Container{
+										{
+											Name: "rabbitmq",
+											LivenessProbe: &corev1.Probe{
+												ProbeHandler: corev1.ProbeHandler{
+													Exec: &corev1.ExecAction{
+														Command: []string{"custom-liveness-probe", "arg1"},
+													},
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					}
+
+					builder = &resource.RabbitmqResourceBuilder{
+						Instance: &instance,
+						Scheme:   scheme,
+					}
+					stsBuilder := builder.StatefulSet()
+					Expect(stsBuilder.Update(statefulSet)).To(Succeed())
+
+					Expect(statefulSet.Spec.Template.Spec.Containers[0].LivenessProbe.ProbeHandler).To(Equal(
+						corev1.ProbeHandler{
+							Exec: &corev1.ExecAction{
+								Command: []string{"custom-liveness-probe", "arg1"},
+							},
+						},
+					))
+
+				})
+
 				Context("TopologySpreadConstraints composition", func() {
 					BeforeEach(func() {
 						instance.Spec.Override.StatefulSet = &rabbitmqv1beta1.StatefulSet{