test: add steps: CreateNetworkPolicies, WaitPodsReady

Signed-off-by: Alex Castilio dos Santos <[email protected]>

test/e2e/framework/kubernetes/wait-pod-ready.go

@@ -10,6 +10,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/clientcmd"
 )
 
 const (
@@ -19,6 +20,53 @@ const (
 	printInterval = 5 // print to stdout every 5 iterations
 )
 
+type WaitPodsReady struct {
+	KubeConfigFilePath string
+	Namespace          string
+	LabelSelector      string
+	DryRun             bool
+	// Name               string
+	// Timeout       time.Duration
+	// AllPods            bool
+}
+
+// Useful when wanting to do parameter checking, for example
+// if a parameter length is known to be required less than 80 characters,
+// do this here so we don't find out later on when we run the step
+// when possible, try to avoid making external calls, this should be fast and simple
+func (w *WaitPodsReady) Prevalidate() error {
+	return nil
+}
+
+// Primary step where test logic is executed
+// Returning an error will cause the test to fail
+func (w *WaitPodsReady) Run() error {
+
+	if w.DryRun {
+		return nil
+	}
+
+	config, err := clientcmd.BuildConfigFromFlags("", w.KubeConfigFilePath)
+	if err != nil {
+		return fmt.Errorf("error building kubeconfig: %w", err)
+	}
+
+	clientset, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		return fmt.Errorf("error creating Kubernetes client: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeoutSeconds*time.Second)
+	defer cancel()
+
+	return WaitForPodReady(ctx, clientset, w.Namespace, w.LabelSelector)
+}
+
+// Require for background steps
+func (w *WaitPodsReady) Stop() error {
+	return nil
+}
+
 func WaitForPodReady(ctx context.Context, clientset *kubernetes.Clientset, namespace, labelSelector string) error {
 	podReadyMap := make(map[string]bool)
 

test/e2e/framework/scaletest/create-network-policies.go

@@ -0,0 +1,169 @@
+package scaletest
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	e2ekubernetes "github.com/microsoft/retina/test/e2e/framework/kubernetes"
+	yaml "gopkg.in/yaml.v2"
+
+	"github.com/microsoft/retina/test/e2e/framework/scaletest/templates"
+	netv1 "k8s.io/api/networking/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/clientcmd"
+)
+
+type CreateNetworkPolicies struct {
+	KubeConfigFilePath          string
+	Namespace                   string
+	NumNetworkPolicies          int
+	NumUnappliedNetworkPolicies int
+	NumSharedLabelsPerPod       int
+	DryRun                      bool
+}
+
+// Useful when wanting to do parameter checking, for example
+// if a parameter length is known to be required less than 80 characters,
+// do this here so we don't find out later on when we run the step
+// when possible, try to avoid making external calls, this should be fast and simple
+func (c *CreateNetworkPolicies) Prevalidate() error {
+	return nil
+}
+
+// Primary step where test logic is executed
+// Returning an error will cause the test to fail
+func (c *CreateNetworkPolicies) Run() error {
+
+	config, err := clientcmd.BuildConfigFromFlags("", c.KubeConfigFilePath)
+	if err != nil {
+		return fmt.Errorf("error building kubeconfig: %w", err)
+	}
+
+	clientset, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		return fmt.Errorf("error creating Kubernetes client: %w", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeoutSeconds*time.Second)
+	defer cancel()
+
+	networkPolicies := c.generateNetworkPolicies(c.NumNetworkPolicies)
+
+	for _, np := range networkPolicies {
+		if c.DryRun {
+			y, err := yaml.Marshal(np)
+			if err != nil {
+				return fmt.Errorf("error marshalling resource: %w", err)
+			}
+			log.Printf("%s", string(y))
+		} else {
+			e2ekubernetes.CreateResource(ctx, np, clientset)
+		}
+	}
+
+	return nil
+}
+
+// Require for background steps
+func (c *CreateNetworkPolicies) Stop() error {
+	return nil
+}
+
+func (c *CreateNetworkPolicies) generateNetworkPolicies(numPolicies int) []runtime.Object {
+	objs := []runtime.Object{}
+	for i := 0; i < numPolicies; i++ {
+		name := fmt.Sprintf("policy-%05d", i)
+
+		template := templates.NetworkPolicy.DeepCopy()
+
+		template.Name = name
+		template.Namespace = c.Namespace
+
+		valNum := i
+		if valNum >= c.NumSharedLabelsPerPod-2 {
+			valNum = c.NumSharedLabelsPerPod - 2
+		}
+
+		template.Spec.PodSelector.MatchLabels[fmt.Sprintf("shared-lab-%05d", valNum)] = "val"
+
+		ingressNum := valNum + 1
+		template.Spec.Ingress = []netv1.NetworkPolicyIngressRule{
+			{
+				From: []netv1.NetworkPolicyPeer{
+					{
+						PodSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								fmt.Sprintf("shared-lab-%05d", ingressNum): "val",
+							},
+						},
+					},
+				},
+			},
+		}
+
+		egressNum := valNum + 2
+		template.Spec.Egress = []netv1.NetworkPolicyEgressRule{
+			{
+				To: []netv1.NetworkPolicyPeer{
+					{
+						PodSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								fmt.Sprintf("shared-lab-%05d", egressNum): "val",
+							},
+						},
+					},
+				},
+			},
+		}
+
+		objs = append(objs, template)
+	}
+
+	for i := 0; i < c.NumUnappliedNetworkPolicies; i++ {
+		name := fmt.Sprintf("unapplied-policy-%05d", i)
+
+		template := templates.NetworkPolicy.DeepCopy()
+
+		template.Name = name
+		template.Namespace = c.Namespace
+
+		template.Spec.PodSelector.MatchLabels["non-existent-key"] = "val"
+
+		template.Spec.Ingress = []netv1.NetworkPolicyIngressRule{
+			{
+				From: []netv1.NetworkPolicyPeer{
+					{
+						PodSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"non-existent-key": "val",
+							},
+						},
+					},
+				},
+			},
+		}
+
+		template.Spec.Egress = []netv1.NetworkPolicyEgressRule{
+			{
+				To: []netv1.NetworkPolicyPeer{
+					{
+						PodSelector: &metav1.LabelSelector{
+							MatchLabels: map[string]string{
+								"non-existent-key": "val",
+							},
+						},
+					},
+				},
+			},
+		}
+
+		objs = append(objs, template)
+
+	}
+
+	return objs
+}

test/e2e/framework/scaletest/templates/networkpolicy.go

@@ -0,0 +1,24 @@
+package templates
+
+import (
+	netv1 "k8s.io/api/networking/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+var (
+	NetworkPolicy = netv1.NetworkPolicy{
+		TypeMeta: metav1.TypeMeta{
+			Kind:       "NetworkPolicy",
+			APIVersion: "networking.k8s.io/v1",
+		},
+		ObjectMeta: metav1.ObjectMeta{
+			Name: "template-network-policy",
+		},
+		Spec: netv1.NetworkPolicySpec{
+			PolicyTypes: []netv1.PolicyType{
+				"Ingress",
+				"Egress",
+			},
+		},
+	}
+)

test/e2e/jobs/scale.go

@@ -84,6 +84,18 @@ func ScaleTest(opt *scaletest.Options) *types.Job {
 		DryRun:                opt.DryRun,
 	}, nil)
 
+	// Apply network policies (applied and unapplied)
+	job.AddStep(&scaletest.CreateNetworkPolicies{
+		NumNetworkPolicies:    opt.NumNetworkPolicies,
+		NumSharedLabelsPerPod: opt.NumSharedLabelsPerPod,
+		DryRun:                opt.DryRun,
+	}, nil)
+
+	job.AddStep(&kubernetes.WaitPodsReady{
+		LabelSelector: "is-real=true",
+		DryRun:        opt.DryRun,
+	}, nil)
+
 	job.AddStep(&scaletest.DeleteAndReAddLabels{
 		DryRun:                opt.DryRun,
 		DeleteLabels:          opt.DeleteLabels,
@@ -92,5 +104,13 @@ func ScaleTest(opt *scaletest.Options) *types.Job {
 		NumSharedLabelsPerPod: opt.NumSharedLabelsPerPod,
 	}, nil)
 
+	// TODO: Add steps to get the state of the cluster
+
+	// job.AddStep(&kubernetes.GetDeployment{})
+
+	// job.AddStep(&kubernetes.GetDaemonSet{})
+
+	// job.AddStep(&kubernetes.DescribePods{})
+
 	return job
 }