From a77f8476d447de61e47bd59a2f4b8b563f70c112 Mon Sep 17 00:00:00 2001
From: Nick Tran <10810510+njtran@users.noreply.github.com>
Date: Fri, 26 Jan 2024 16:12:55 -0800
Subject: [PATCH] test: add command and prestop into test pod options (#973)

---
 pkg/test/pods.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/pkg/test/pods.go b/pkg/test/pods.go
index 52e56ec5bd..6bf11173dd 100644
--- a/pkg/test/pods.go
+++ b/pkg/test/pods.go
@@ -55,6 +55,8 @@ type PodOptions struct {
 	TerminationGracePeriodSeconds *int64
 	ReadinessProbe                *v1.Probe
 	LivenessProbe                 *v1.Probe
+	PreStopSleep                  *int64
+	Command                       []string
 }
 
 type PDBOptions struct {
@@ -143,6 +145,25 @@ func Pod(overrides ...PodOptions) *v1.Pod {
 			Phase:      options.Phase,
 		},
 	}
+	// If PreStopSleep is enabled, add it to each of the containers.
+	// Can't use v1.LifecycleHandler == v1.SleepAction as that's a feature gate in Alpha 1.29.
+	// https://kubernetes.io/docs/concepts/containers/container-lifecycle-hooks/#hook-handler-implementations
+	if options.PreStopSleep != nil {
+		p.Spec.Containers[0].Lifecycle = &v1.Lifecycle{
+			PreStop: &v1.LifecycleHandler{
+				Exec: &v1.ExecAction{
+					Command: []string{
+						"/bin/sh",
+						"-c",
+						fmt.Sprintf("sleep %d", *options.PreStopSleep),
+					},
+				},
+			},
+		}
+	}
+	if options.Command != nil {
+		p.Spec.Containers[0].Command = options.Command
+	}
 	if options.InitImage != "" {
 		p.Spec.InitContainers = []v1.Container{{
 			Name:      RandomName(),