Use int or percentage in pod termination count

Signed-off-by: Pablo Chacin <[email protected]>

e2e/disruptors/pod_e2e_test.go

@@ -314,7 +314,7 @@ func Test_PodDisruptor(t *testing.T) {
 		}
 
 		fault := disruptors.TerminatePodsFault{
-			Count:   1,
+			Count:   intstr.FromInt32(1),
 			Timeout: 10 * time.Second,
 		}
 
@@ -323,8 +323,8 @@ func Test_PodDisruptor(t *testing.T) {
 			t.Fatalf("terminating pods: %v", err)
 		}
 
-		if len(terminated) != fault.Count {
-			t.Fatalf("Invalid number of pods deleted. Expected %d got %d", fault.Count, len(terminated))
+		if len(terminated) != int(fault.Count.Int32()) {
+			t.Fatalf("Invalid number of pods deleted. Expected %d got %d", fault.Count.Int32(), len(terminated))
 		}
 
 		for _, pod := range terminated {

e2e/disruptors/service_e2e_test.go

@@ -186,7 +186,7 @@ func Test_ServiceDisruptor(t *testing.T) {
 		}
 
 		fault := disruptors.TerminatePodsFault{
-			Count:   1,
+			Count:   intstr.FromInt32(1),
 			Timeout: 10 * time.Second,
 		}
 
@@ -195,8 +195,8 @@ func Test_ServiceDisruptor(t *testing.T) {
 			t.Fatalf("terminating pods: %v", err)
 		}
 
-		if len(terminated) != fault.Count {
-			t.Fatalf("Invalid number of pods deleted. Expected %d got %d", fault.Count, len(terminated))
+		if len(terminated) != int(fault.Count.Int32()) {
+			t.Fatalf("Invalid number of pods deleted. Expected %d got %d", fault.Count.Int32(), len(terminated))
 		}
 
 		for _, pod := range terminated {

pkg/disruptors/terminate.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/grafana/xk6-disruptor/pkg/kubernetes/helpers"
+	"github.com/grafana/xk6-disruptor/pkg/types/intstr"
 	corev1 "k8s.io/api/core/v1"
 )
 
@@ -30,8 +31,8 @@ type PodFaultInjector interface {
 
 // TerminatePodsFault specifies a fault that will terminate a set of pods
 type TerminatePodsFault struct {
-	// Count indicates how many pods to terminate
-	Count int
+	// Count indicates how many pods to terminate. Can be a number or a percentage or targets
+	Count intstr.IntOrString
 	// Timeout specifies the maximum time to wait for a pod to terminate
 	Timeout time.Duration
 }

pkg/utils/kubernetes.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"fmt"
+	"math"
 
 	"github.com/grafana/xk6-disruptor/pkg/types/intstr"
 	corev1 "k8s.io/api/core/v1"
@@ -77,10 +78,29 @@ func PodNames(pods []corev1.Pod) []string {
 	return names
 }
 
-// Sample a subset of the given list of Pods
-func Sample(pods []corev1.Pod, count int) ([]corev1.Pod, error) {
-	if count > len(pods) {
+// Sample a subset of the given list of Pods. The count is defined as a int or a string representing a percentage.
+// If the count is a percentage and there are no enough elements in the pod list, the number is rounded up.
+// If the list is not empty, at least one element is returned
+// For example 25% of a list of 2 pods will return one pod.
+func Sample(pods []corev1.Pod, count intstr.IntOrString) ([]corev1.Pod, error) {
+	var sampleSize int
+	if count.IsInt() {
+		sampleSize = int(count.Int32())
+	} else {
+		percentage, ok := count.AsPercentage()
+		if !ok {
+			return nil, fmt.Errorf("count is not a valid percentage: %s", count)
+		}
+
+		if percentage == 0 {
+			return nil, nil
+		}
+
+		sampleSize = int(math.Max(1, math.Round(float64(len(pods)*int(percentage))/100)))
+	}
+
+	if sampleSize > len(pods) {
 		return nil, fmt.Errorf("not enough elements to sample")
 	}
-	return pods[:count], nil
+	return pods[:sampleSize], nil
 }

pkg/utils/kubernetes_test.go

@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"fmt"
 	"testing"
 
 	corev1 "k8s.io/api/core/v1"
@@ -182,3 +183,104 @@ func Test_GetTargetPort(t *testing.T) {
 		})
 	}
 }
+
+func Test_Sample(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		title       string
+		numPods     int
+		count       intstr.IntOrString
+		expect      int
+		expectError bool
+	}{
+		{
+			title:       "select one pod",
+			numPods:     3,
+			count:       intstr.FromInt32(1),
+			expect:      1,
+			expectError: false,
+		},
+		{
+			title:       "select all pods",
+			numPods:     3,
+			count:       intstr.FromInt32(3),
+			expect:      3,
+			expectError: false,
+		},
+		{
+			title:       "select too many pods",
+			numPods:     3,
+			count:       intstr.FromInt32(4),
+			expect:      0,
+			expectError: true,
+		},
+		{
+			title:       "select 25% of pods",
+			numPods:     3,
+			count:       intstr.FromString("25%"),
+			expect:      1,
+			expectError: false,
+		},
+		{
+			title:       "select 30% of pods",
+			numPods:     3,
+			count:       intstr.FromString("30%"),
+			expect:      1,
+			expectError: false,
+		},
+		{
+			title:       "select 50% of pods",
+			numPods:     3,
+			count:       intstr.FromString("50%"),
+			expect:      2,
+			expectError: false,
+		},
+		{
+			title:       "select 100% of pods",
+			numPods:     3,
+			count:       intstr.FromString("100%"),
+			expect:      3,
+			expectError: false,
+		},
+		{
+			title:       "select 0% of pods",
+			numPods:     3,
+			count:       intstr.FromString("0%"),
+			expect:      0,
+			expectError: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.title, func(t *testing.T) {
+			t.Parallel()
+
+			pods := []corev1.Pod{}
+			for i := 0; i < tc.numPods; i++ {
+				podName := fmt.Sprintf("pod-%d", i)
+				pods = append(pods, builders.NewPodBuilder(podName).Build())
+			}
+
+			sample, err := Sample(pods, tc.count)
+
+			if err != nil && !tc.expectError {
+				t.Fatalf("failed %v", err)
+			}
+
+			if err == nil && tc.expectError {
+				t.Fatalf("should had failed")
+			}
+
+			if err != nil && tc.expectError {
+				return
+			}
+
+			if len(sample) != tc.expect {
+				t.Fatalf("expected %d pods got %d", tc.expect, len(sample))
+			}
+		})
+	}
+}