From d2a1fa9e7d25d84cda9989bb5db35adb34d9c5b3 Mon Sep 17 00:00:00 2001
From: Eduardo Espinoza Perez <esespinozap@falabella.cl>
Date: Mon, 21 Nov 2022 13:54:54 -0300
Subject: [PATCH] feat(k8s): add networkpolicy support

---
 modules/k8s/networkpolicy.go | 53 ++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 modules/k8s/networkpolicy.go

diff --git a/modules/k8s/networkpolicy.go b/modules/k8s/networkpolicy.go
new file mode 100644
index 000000000..124028e6e
--- /dev/null
+++ b/modules/k8s/networkpolicy.go
@@ -0,0 +1,53 @@
+package k8s
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/gruntwork-io/terratest/modules/logger"
+	"github.com/gruntwork-io/terratest/modules/retry"
+	"github.com/gruntwork-io/terratest/modules/testing"
+	"github.com/stretchr/testify/require"
+	networkingv1 "k8s.io/api/networking/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// GetNetworkPolicy returns a Kubernetes networkpolicy resource in the provided namespace with the given name. The namespace used
+// is the one provided in the KubectlOptions. This will fail the test if there is an error.
+func GetNetworkPolicy(t testing.TestingT, options *KubectlOptions, networkPolicyName string) *networkingv1.NetworkPolicy {
+	networkPolicy, err := GetNetworkPolicyE(t, options, networkPolicyName)
+	require.NoError(t, err)
+	return networkPolicy
+}
+
+// GetNetworkPolicyE returns a Kubernetes networkpolicy resource in the provided namespace with the given name. The namespace used
+// is the one provided in the KubectlOptions.
+func GetNetworkPolicyE(t testing.TestingT, options *KubectlOptions, networkPolicyName string) (*networkingv1.NetworkPolicy, error) {
+	clientset, err := GetKubernetesClientFromOptionsE(t, options)
+	if err != nil {
+		return nil, err
+	}
+	return clientset.NetworkingV1().NetworkPolicies(options.Namespace).Get(context.Background(), networkPolicyName, metav1.GetOptions{})
+}
+
+// WaitUntilNetworkPolicyAvailable waits until the networkpolicy is present on the cluster in cases where it is not immediately
+// available (for example, when using ClusterIssuer to request a certificate).
+func WaitUntilNetworkPolicyAvailable(t testing.TestingT, options *KubectlOptions, networkPolicyName string, retries int, sleepBetweenRetries time.Duration) {
+	statusMsg := fmt.Sprintf("Wait for networkpolicy %s to be provisioned.", networkPolicyName)
+	message := retry.DoWithRetry(
+		t,
+		statusMsg,
+		retries,
+		sleepBetweenRetries,
+		func() (string, error) {
+			_, err := GetNetworkPolicyE(t, options, networkPolicyName)
+			if err != nil {
+				return "", err
+			}
+
+			return "networkpolicy is now available", nil
+		},
+	)
+	logger.Logf(t, message)
+}