From b14b4686619683c354fdf19cdf4e5bcff88965b9 Mon Sep 17 00:00:00 2001
From: Marek Siarkowicz <siarkowicz@google.com>
Date: Mon, 8 May 2023 19:58:32 +0200
Subject: [PATCH] tests/robustness: Make weighted pick random generic

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>
---
 tests/robustness/linearizability_test.go | 38 ++++++++++++------------
 tests/robustness/traffic.go              | 14 ++++-----
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/tests/robustness/linearizability_test.go b/tests/robustness/linearizability_test.go
index 5a441cdf4..5ecc9e2f2 100644
--- a/tests/robustness/linearizability_test.go
+++ b/tests/robustness/linearizability_test.go
@@ -45,14 +45,14 @@ var (
 			keyCount:     10,
 			leaseTTL:     DefaultLeaseTTL,
 			largePutSize: 32769,
-			writeChoices: []choiceWeight{
-				{choice: string(Put), weight: 45},
-				{choice: string(LargePut), weight: 5},
-				{choice: string(Delete), weight: 10},
-				{choice: string(MultiOpTxn), weight: 10},
-				{choice: string(PutWithLease), weight: 10},
-				{choice: string(LeaseRevoke), weight: 10},
-				{choice: string(CompareAndSet), weight: 10},
+			writeChoices: []choiceWeight[etcdRequestType]{
+				{choice: Put, weight: 45},
+				{choice: LargePut, weight: 5},
+				{choice: Delete, weight: 10},
+				{choice: MultiOpTxn, weight: 10},
+				{choice: PutWithLease, weight: 10},
+				{choice: LeaseRevoke, weight: 10},
+				{choice: CompareAndSet, weight: 10},
 			},
 		},
 	}
@@ -66,10 +66,10 @@ var (
 			keyCount:     10,
 			largePutSize: 32769,
 			leaseTTL:     DefaultLeaseTTL,
-			writeChoices: []choiceWeight{
-				{choice: string(Put), weight: 85},
-				{choice: string(MultiOpTxn), weight: 10},
-				{choice: string(LargePut), weight: 5},
+			writeChoices: []choiceWeight[etcdRequestType]{
+				{choice: Put, weight: 85},
+				{choice: MultiOpTxn, weight: 10},
+				{choice: LargePut, weight: 5},
 			},
 		},
 	}
@@ -82,10 +82,10 @@ var (
 			averageKeyCount: 5,
 			resource:        "pods",
 			namespace:       "default",
-			writeChoices: []choiceWeight{
-				{choice: string(KubernetesUpdate), weight: 75},
-				{choice: string(KubernetesDelete), weight: 15},
-				{choice: string(KubernetesCreate), weight: 10},
+			writeChoices: []choiceWeight[KubernetesRequestType]{
+				{choice: KubernetesUpdate, weight: 75},
+				{choice: KubernetesDelete, weight: 15},
+				{choice: KubernetesCreate, weight: 10},
 			},
 		},
 	}
@@ -99,9 +99,9 @@ var (
 			keyCount:     10,
 			largePutSize: 8196,
 			leaseTTL:     DefaultLeaseTTL,
-			writeChoices: []choiceWeight{
-				{choice: string(Put), weight: 95},
-				{choice: string(LargePut), weight: 5},
+			writeChoices: []choiceWeight[etcdRequestType]{
+				{choice: Put, weight: 95},
+				{choice: LargePut, weight: 5},
 			},
 		},
 	}
diff --git a/tests/robustness/traffic.go b/tests/robustness/traffic.go
index a3b61133e..9fb5335ae 100644
--- a/tests/robustness/traffic.go
+++ b/tests/robustness/traffic.go
@@ -110,7 +110,7 @@ type Traffic interface {
 
 type etcdTraffic struct {
 	keyCount     int
-	writeChoices []choiceWeight
+	writeChoices []choiceWeight[etcdRequestType]
 	leaseTTL     int64
 	largePutSize int
 }
@@ -132,7 +132,7 @@ type kubernetesTraffic struct {
 	averageKeyCount int
 	resource        string
 	namespace       string
-	writeChoices    []choiceWeight
+	writeChoices    []choiceWeight[KubernetesRequestType]
 }
 
 type KubernetesRequestType string
@@ -174,7 +174,7 @@ func (t kubernetesTraffic) Write(ctx context.Context, c *recordingClient, ids id
 		if len(objects) > t.averageKeyCount*3/2 {
 			err = t.Delete(writeCtx, c, string(randomPod.Key), randomPod.ModRevision)
 		} else {
-			op := KubernetesRequestType(pickRandom(t.writeChoices))
+			op := pickRandom(t.writeChoices)
 			switch op {
 			case KubernetesDelete:
 				err = t.Delete(writeCtx, c, string(randomPod.Key), randomPod.ModRevision)
@@ -256,7 +256,7 @@ func (t etcdTraffic) Write(ctx context.Context, c *recordingClient, limiter *rat
 	writeCtx, cancel := context.WithTimeout(ctx, RequestTimeout)
 
 	var err error
-	switch etcdRequestType(pickRandom(t.writeChoices)) {
+	switch pickRandom(t.writeChoices) {
 	case Put:
 		err = c.Put(writeCtx, key, fmt.Sprintf("%d", id.RequestId()))
 	case LargePut:
@@ -356,12 +356,12 @@ func randString(size int) string {
 	return data.String()
 }
 
-type choiceWeight struct {
-	choice string
+type choiceWeight[T any] struct {
+	choice T
 	weight int
 }
 
-func pickRandom(choices []choiceWeight) string {
+func pickRandom[T any](choices []choiceWeight[T]) T {
 	sum := 0
 	for _, op := range choices {
 		sum += op.weight