From 49887031953b91e4e00b56b34e9869270995307e Mon Sep 17 00:00:00 2001
From: Dejan Pejchev <pejcev.dejan@gmail.com>
Date: Sun, 14 Apr 2024 00:54:51 +0200
Subject: [PATCH] add e2e test for ttl seconds after finished in jobset

---
 pkg/util/testing/wrappers.go |  5 ++++
 test/e2e/e2e_test.go         | 47 ++++++++++++++++++++++++++++++++++++
 test/util/util.go            | 37 ++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/pkg/util/testing/wrappers.go b/pkg/util/testing/wrappers.go
index 4f3d4bb6b..03a1ec38b 100644
--- a/pkg/util/testing/wrappers.go
+++ b/pkg/util/testing/wrappers.go
@@ -153,6 +153,11 @@ func (j *JobSetWrapper) DeletionTimestamp(deletionTimestamp *metav1.Time) *JobSe
 
 }
 
+func (j *JobSetWrapper) Finalizers(finalizers []string) *JobSetWrapper {
+	j.ObjectMeta.Finalizers = finalizers
+	return j
+}
+
 // ReplicatedJobWrapper wraps a ReplicatedJob.
 type ReplicatedJobWrapper struct {
 	jobset.ReplicatedJob
diff --git a/test/e2e/e2e_test.go b/test/e2e/e2e_test.go
index e31064042..9629e4957 100644
--- a/test/e2e/e2e_test.go
+++ b/test/e2e/e2e_test.go
@@ -105,6 +105,31 @@ var _ = ginkgo.Describe("JobSet", func() {
 			util.JobSetCompleted(ctx, k8sClient, js, timeout)
 		})
 	})
+	ginkgo.When("ttl seconds after finished is set", func() {
+		ginkgo.It("should clean up the completed jobset after configured ttl seconds expire", func() {
+			ctx := context.Background()
+
+			// Create JobSet.
+			testFinalizer := "fake.example.com/blockDeletion"
+			ginkgo.By("creating jobset with ttl seconds after finished")
+			js := sleepTestJobSet(ns).Finalizers([]string{testFinalizer}).TTLSecondsAfterFinished(5).Obj()
+
+			// Verify jobset created successfully.
+			ginkgo.By("checking that jobset creation succeeds")
+			gomega.Expect(k8sClient.Create(ctx, js)).Should(gomega.Succeed())
+
+			// Check jobset status if specified.
+			ginkgo.By("checking jobset condition")
+			util.JobSetCompleted(ctx, k8sClient, js, timeout)
+
+			// We remove the jobset finalizer, so it can get deleted when ttl expires.
+			util.RemoveJobSetFinalizer(ctx, k8sClient, js, testFinalizer, timeout)
+
+			// Check jobset is cleaned up after ttl seconds.
+			ginkgo.By("checking jobset is cleaned up after ttl seconds")
+			util.JobSetDeleted(ctx, k8sClient, js, timeout)
+		})
+	})
 
 }) // end of Describe
 
@@ -197,3 +222,25 @@ func pingTestJobSetSubdomain(ns *corev1.Namespace) *testing.JobSetWrapper {
 			Replicas(int32(replicas)).
 			Obj())
 }
+
+func sleepTestJobSet(ns *corev1.Namespace) *testing.JobSetWrapper {
+	jsName := "js"
+	rjobName := "rjob"
+	replicas := 4
+	return testing.MakeJobSet(jsName, ns.Name).
+		ReplicatedJob(testing.MakeReplicatedJob(rjobName).
+			Job(testing.MakeJobTemplate("job", ns.Name).
+				PodSpec(corev1.PodSpec{
+					RestartPolicy: "Never",
+					Containers: []corev1.Container{
+						{
+							Name:    "sleep-test-container",
+							Image:   "bash:latest",
+							Command: []string{"bash", "-c"},
+							Args:    []string{"sleep 20"},
+						},
+					},
+				}).Obj()).
+			Replicas(int32(replicas)).
+			Obj())
+}
diff --git a/test/util/util.go b/test/util/util.go
index 8114f6b6e..8a10b1da4 100644
--- a/test/util/util.go
+++ b/test/util/util.go
@@ -188,3 +188,40 @@ func ExpectJobsDeletionTimestamp(ctx context.Context, c client.Client, js *jobse
 		return numJobs == 0, nil
 	}, timeout, interval).Should(gomega.Equal(true))
 }
+
+func JobSetDeleted(ctx context.Context, k8sClient client.Client, js *jobset.JobSet, timeout time.Duration) {
+	ginkgo.By("checking jobset is deleted")
+	gomega.Eventually(func() (bool, error) {
+		err := k8sClient.Get(ctx, types.NamespacedName{Namespace: js.Namespace, Name: js.Name}, js)
+		if apierrors.IsNotFound(err) {
+			return true, nil
+		}
+		return false, err
+	}, timeout, interval).Should(gomega.Equal(true))
+}
+
+// RemoveJobSetFinalizer removes the provided finalizer from the jobset and updates it.
+func RemoveJobSetFinalizer(ctx context.Context, k8sClient client.Client, js *jobset.JobSet, finalizer string, timeout time.Duration) {
+	ginkgo.By("removing jobset finalizers")
+	gomega.Eventually(func() (bool, error) {
+		// We get the latest version of the jobset before removing the finalizer.
+		var fresh jobset.JobSet
+		if err := k8sClient.Get(ctx, types.NamespacedName{Name: js.Name, Namespace: js.Namespace}, &fresh); err != nil {
+			return false, err
+		}
+		removeJobSetFinalizer(&fresh, finalizer)
+		if err := k8sClient.Update(ctx, &fresh); err != nil {
+			return false, err
+		}
+		return true, nil
+	}, timeout, interval).Should(gomega.Equal(true))
+}
+
+// removeJobSetFinalizer removes the provided finalizer from the jobset.
+func removeJobSetFinalizer(js *jobset.JobSet, finalizer string) {
+	for i, f := range js.Finalizers {
+		if f == finalizer {
+			js.Finalizers = append(js.Finalizers[:i], js.Finalizers[i+1:]...)
+		}
+	}
+}