From 3af09ab04a671d79379dc19178842b82ded7e973 Mon Sep 17 00:00:00 2001
From: Matthew Cary <mattcary@google.com>
Date: Wed, 15 Jun 2022 13:33:02 -0700
Subject: [PATCH] add TestAutodeleteOwnerRefs statefulset integration test

Change-Id: I3aa3220df59c33fc40f1d8e08e055caa1f6fb4a3
---
 .../statefulset/statefulset_test.go           | 86 +++++++++++++++++++
 test/integration/statefulset/util.go          | 33 ++++++-
 2 files changed, 117 insertions(+), 2 deletions(-)

diff --git a/test/integration/statefulset/statefulset_test.go b/test/integration/statefulset/statefulset_test.go
index 2e2d99c14c3..5c26be3f45b 100644
--- a/test/integration/statefulset/statefulset_test.go
+++ b/test/integration/statefulset/statefulset_test.go
@@ -403,3 +403,89 @@ func TestStatefulSetStatusWithPodFail(t *testing.T) {
 		t.Fatalf("StatefulSet %s status has %d replicas, want replicas %d: %v", sts.Name, gotReplicas, wantReplicas, err)
 	}
 }
+
+func TestAutodeleteOwnerRefs(t *testing.T) {
+	tests := []struct {
+		name              string
+		policy            appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy
+		expectPodOwnerRef bool
+		expectSetOwnerRef bool
+	}{
+		{
+			name: "always retain",
+			policy: appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenDeleted: appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+				WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+			},
+			expectPodOwnerRef: false,
+			expectSetOwnerRef: false,
+		},
+		{
+			name: "delete on scaledown only",
+			policy: appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenDeleted: appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+				WhenScaled:  appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+			},
+			expectPodOwnerRef: true,
+			expectSetOwnerRef: false,
+		},
+		{
+			name: "delete with set only",
+			policy: appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+				WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+			},
+			expectPodOwnerRef: false,
+			expectSetOwnerRef: true,
+		},
+		{
+			name: "always delete",
+			policy: appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+				WhenScaled:  appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+			},
+			expectPodOwnerRef: true,
+			expectSetOwnerRef: true,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.StatefulSetAutoDeletePVC, true)()
+			closeFn, rm, informers, c := scSetup(t)
+			defer closeFn()
+			ns := framework.CreateNamespaceOrDie(c, "test-autodelete-ownerrefs", t)
+			defer framework.DeleteNamespaceOrDie(c, ns, t)
+			cancel := runControllerAndInformers(rm, informers)
+			defer cancel()
+
+			sts := newSTS("sts", ns.Name, 3)
+			sts.Spec.PersistentVolumeClaimRetentionPolicy = &test.policy
+			stss, _ := createSTSsPods(t, c, []*appsv1.StatefulSet{sts}, []*v1.Pod{})
+			sts = stss[0]
+			waitSTSStable(t, c, sts)
+
+			// Verify StatefulSet ownerref has been added as appropriate.
+			pvcClient := c.CoreV1().PersistentVolumeClaims(ns.Name)
+			pvcs := getStatefulSetPVCs(t, pvcClient, sts)
+			for _, pvc := range pvcs {
+				verifyOwnerRef(t, pvc, "StatefulSet", test.expectSetOwnerRef)
+				verifyOwnerRef(t, pvc, "Pod", false)
+			}
+
+			// Scale down to 1 pod and verify Pod ownerrefs as appropriate.
+			one := int32(1)
+			sts.Spec.Replicas = &one
+			waitSTSStable(t, c, sts)
+
+			pvcs = getStatefulSetPVCs(t, pvcClient, sts)
+			for i, pvc := range pvcs {
+				verifyOwnerRef(t, pvc, "StatefulSet", test.expectSetOwnerRef)
+				if i == 0 {
+					verifyOwnerRef(t, pvc, "Pod", false)
+				} else {
+					verifyOwnerRef(t, pvc, "Pod", test.expectPodOwnerRef)
+				}
+			}
+		})
+	}
+}
diff --git a/test/integration/statefulset/util.go b/test/integration/statefulset/util.go
index 85ef679439c..2090e977f41 100644
--- a/test/integration/statefulset/util.go
+++ b/test/integration/statefulset/util.go
@@ -110,8 +110,8 @@ func newSTS(name, namespace string, replicas int) *appsv1.StatefulSet {
 						{
 							Name: "datadir",
 							VolumeSource: v1.VolumeSource{
-								HostPath: &v1.HostPathVolumeSource{
-									Path: fmt.Sprintf("/tmp/%v", "datadir"),
+								PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{
+									ClaimName: "fake-pvc-name",
 								},
 							},
 						},
@@ -279,6 +279,35 @@ func getPods(t *testing.T, podClient typedv1.PodInterface, labelMap map[string]s
 	return pods
 }
 
+func getStatefulSetPVCs(t *testing.T, pvcClient typedv1.PersistentVolumeClaimInterface, sts *appsv1.StatefulSet) []*v1.PersistentVolumeClaim {
+	pvcs := []*v1.PersistentVolumeClaim{}
+	for i := int32(0); i < *sts.Spec.Replicas; i++ {
+		pvcName := fmt.Sprintf("%s-%s-%d", sts.Spec.VolumeClaimTemplates[0].Name, sts.Name, i)
+		pvc, err := pvcClient.Get(context.TODO(), pvcName, metav1.GetOptions{})
+		if err != nil {
+			t.Fatalf("failed to get PVC %s: %v", pvcName, err)
+		}
+		pvcs = append(pvcs, pvc)
+	}
+	return pvcs
+}
+
+func verifyOwnerRef(t *testing.T, pvc *v1.PersistentVolumeClaim, kind string, expected bool) {
+	found := false
+	for _, ref := range pvc.GetOwnerReferences() {
+		if ref.Kind == kind {
+			if expected {
+				found = true
+			} else {
+				t.Fatalf("Found %s ref but expected none for PVC %s", kind, pvc.Name)
+			}
+		}
+	}
+	if expected && !found {
+		t.Fatalf("Expected %s ref but found none for PVC %s", kind, pvc.Name)
+	}
+}
+
 func updateSTS(t *testing.T, stsClient typedappsv1.StatefulSetInterface, stsName string, updateFunc func(*appsv1.StatefulSet)) *appsv1.StatefulSet {
 	var sts *appsv1.StatefulSet
 	if err := retry.RetryOnConflict(retry.DefaultBackoff, func() error {