Add integration test for Pod delete

Author: Supreeth095

test/framework/manifest/statefulset.go

@@ -0,0 +1,89 @@
+package manifest
+
+import (
+    "github.com/aws/aws-sdk-go-v2/aws"
+    appsv1 "k8s.io/api/apps/v1"
+    corev1 "k8s.io/api/core/v1"
+    metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+type StatefulSetBuilder struct {
+    namespace              string
+    name                   string
+    serviceName            string
+    replicas               int32
+    labels                 map[string]string
+    container             corev1.Container
+    nodeSelector          map[string]string
+    terminationGracePeriod int64
+}
+
+func NewDefaultStatefulSetBuilder() *StatefulSetBuilder {
+    return &StatefulSetBuilder{
+        namespace:              "default",
+        name:                   "statefulset",
+        serviceName:           "statefulset",
+        replicas:               2,
+        labels:                 map[string]string{},
+        container:             NewBusyBoxContainerBuilder().Build(),
+        nodeSelector:          map[string]string{"kubernetes.io/os": "linux"},
+        terminationGracePeriod: 0,
+    }
+}
+
+func (s *StatefulSetBuilder) Namespace(namespace string) *StatefulSetBuilder {
+    s.namespace = namespace
+    return s
+}
+
+func (s *StatefulSetBuilder) Name(name string) *StatefulSetBuilder {
+    s.name = name
+    s.serviceName = name // typically same as name
+    return s
+}
+
+func (s *StatefulSetBuilder) ServiceName(serviceName string) *StatefulSetBuilder {
+    s.serviceName = serviceName
+    return s
+}
+
+func (s *StatefulSetBuilder) Replicas(replicas int32) *StatefulSetBuilder {
+    s.replicas = replicas
+    return s
+}
+
+func (s *StatefulSetBuilder) PodLabel(key, value string) *StatefulSetBuilder {
+    s.labels[key] = value
+    return s
+}
+
+func (s *StatefulSetBuilder) Container(container corev1.Container) *StatefulSetBuilder {
+    s.container = container
+    return s
+}
+
+func (s *StatefulSetBuilder) Build() *appsv1.StatefulSet {
+    return &appsv1.StatefulSet{
+        ObjectMeta: metav1.ObjectMeta{
+            Name:      s.name,
+            Namespace: s.namespace,
+        },
+        Spec: appsv1.StatefulSetSpec{
+            ServiceName: s.serviceName,
+            Replicas:    aws.Int32(s.replicas),
+            Selector: &metav1.LabelSelector{
+                MatchLabels: s.labels,
+            },
+            Template: corev1.PodTemplateSpec{
+                ObjectMeta: metav1.ObjectMeta{
+                    Labels: s.labels,
+                },
+                Spec: corev1.PodSpec{
+                    Containers:                    []corev1.Container{s.container},
+                    NodeSelector:                  s.nodeSelector,
+                    TerminationGracePeriodSeconds: aws.Int64(s.terminationGracePeriod),
+                },
+            },
+        },
+    }
+}

test/integration/perpodsg/perpodsg_test.go

@@ -28,6 +28,7 @@ import (
 	sgpWrapper "github.com/aws/amazon-vpc-resource-controller-k8s/test/framework/resource/k8s/sgp"
 	"github.com/aws/amazon-vpc-resource-controller-k8s/test/framework/utils"
 	"github.com/samber/lo"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -513,8 +514,108 @@ var _ = Describe("Branch ENI Pods", func() {
 			})
 		})
 	})
+
+	Describe("Test Network Connectivity on Delete and Recreation of Pod", func() {
+		Context("creating statefulset with network connectivity", func() {
+			var resourceMap map[v1.ResourceName]resource.Quantity
+			var statefulSet *appsv1.StatefulSet
+			var container v1.Container
+
+			BeforeEach(func() {
+				resourceMap = map[v1.ResourceName]resource.Quantity{
+					config.ResourceNamePodENI: resource.MustParse("1"),
+				}
+			})
+
+			JustBeforeEach(func() {
+				container = manifest.NewBusyBoxContainerBuilder().
+							Resources(v1.ResourceRequirements{
+								Limits:   resourceMap,
+								Requests: resourceMap,
+							}).
+							Command([]string{ "/bin/sh", "-c",
+									"while true; do if ping -c 1 google.com; then echo 'Successfully pinged google.com'; else echo 'Failed to ping google.com'; exit 1; fi; sleep 30; done",
+									}).
+							Name("network-test").
+							Image("busybox").							
+							Build()
+
+				statefulSet = manifest.NewDefaultStatefulSetBuilder().
+					Namespace(namespace).
+					Name("network-test").
+					PodLabel(podLabelKey, podLabelValue).
+					Container(container).
+					Build()
+			})
+
+			JustAfterEach(func() {
+				By("deleting the statefulset")
+				err = frameWork.K8sClient.Delete(ctx, statefulSet)
+				Expect(err).ToNot(HaveOccurred())
+			})
+
+			Context("when statefulset is created with network connectivity requirements", func() {
+				It("should have all pods running with network access", func() {
+					By("creating security group policy")
+					sgpWrapper.CreateSecurityGroupPolicy(frameWork.K8sClient, ctx, securityGroupPolicy)
+
+					By("creating statefulset")
+					err = frameWork.K8sClient.Create(ctx, statefulSet)
+					Expect(err).ToNot(HaveOccurred())
+
+					By("waiting for statefulset pods to be ready")
+					Eventually(func() bool {
+						err := frameWork.K8sClient.Get(ctx, client.ObjectKey{
+							Namespace: namespace,
+							Name:      statefulSet.Name,
+						}, statefulSet)
+						if err != nil {
+							return false
+						}
+						return statefulSet.Status.ReadyReplicas == *statefulSet.Spec.Replicas
+					}, 300*time.Second, 5*time.Second).Should(BeTrue())
+
+					By("verifying network connectivity of all pods")
+					verify.VerifyNetworkingOfAllPodUsingENI(namespace, podLabelKey, podLabelValue,
+						securityGroups)
+
+					By("force deleting one pod to verify recreation")
+					pods := &v1.PodList{}
+					err = frameWork.K8sClient.List(ctx, pods, client.InNamespace(namespace),
+						client.MatchingLabels(map[string]string{podLabelKey: podLabelValue}))
+					Expect(err).ToNot(HaveOccurred())
+					Expect(pods.Items).ToNot(BeEmpty())
+
+					// Force delete the first pod
+					err = frameWork.K8sClient.Delete(ctx, &pods.Items[0], client.GracePeriodSeconds(0))
+					Expect(err).ToNot(HaveOccurred())
+
+					By("waiting for pod to be recreated")
+					Eventually(func() bool {
+						err := frameWork.K8sClient.Get(ctx, client.ObjectKey{
+							Namespace: namespace,
+							Name:      statefulSet.Name,
+						}, statefulSet)
+						if err != nil {
+							return false
+						}
+						return statefulSet.Status.ReadyReplicas == *statefulSet.Spec.Replicas
+					}, 300*time.Second, 5*time.Second).Should(BeTrue())
+
+					By("verifying network connectivity after pod recreation")
+					verify.VerifyNetworkingOfAllPodUsingENI(namespace, podLabelKey, podLabelValue,
+						securityGroups)
+				})
+			})
+		})
+	})
 })
 
+// Helper functions
+func int32Ptr(i int32) *int32 {
+    return &i
+}
+
 func CreateServiceAccount(serviceAccount *v1.ServiceAccount) {
 	By("create a service account")
 	err := frameWork.K8sClient.Create(ctx, serviceAccount)