From 0df421b37f82ffe543e63e7460b5c64e4c56a203 Mon Sep 17 00:00:00 2001
From: Sam Kleinman <garen@tychoish.com>
Date: Fri, 27 Aug 2021 12:32:40 -0400
Subject: [PATCH] e2e: add weighted random configuration selector (#6869)

When revwing #6807 I assumed that `probSetChoice` worked this way.

I think that the coverage of various configuration options should
generally track what we expect the actual useage to be to focus the
most test coverage on the configurations that are the most prevelent.
---
 go.mod                         |  1 +
 go.sum                         |  2 ++
 test/e2e/generator/generate.go | 12 +++++++++---
 test/e2e/generator/random.go   | 18 ++++++++++++++++++
 4 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index c7723b096..f88b1d689 100644
--- a/go.mod
+++ b/go.mod
@@ -22,6 +22,7 @@ require (
 	github.com/lib/pq v1.10.2
 	github.com/libp2p/go-buffer-pool v0.0.2
 	github.com/minio/highwayhash v1.0.2
+	github.com/mroth/weightedrand v0.4.1 // indirect
 	github.com/oasisprotocol/curve25519-voi v0.0.0-20210609091139-0a56a4bca00b
 	github.com/ory/dockertest v3.3.5+incompatible
 	github.com/prometheus/client_golang v1.11.0
diff --git a/go.sum b/go.sum
index fff11dad3..606ea9f0b 100644
--- a/go.sum
+++ b/go.sum
@@ -636,6 +636,8 @@ github.com/moricho/tparallel v0.2.1 h1:95FytivzT6rYzdJLdtfn6m1bfFJylOJK41+lgv/EH
 github.com/moricho/tparallel v0.2.1/go.mod h1:fXEIZxG2vdfl0ZF8b42f5a78EhjjD5mX8qUplsoSU4k=
 github.com/mozilla/scribe v0.0.0-20180711195314-fb71baf557c1/go.mod h1:FIczTrinKo8VaLxe6PWTPEXRXDIHz2QAwiaBaP5/4a8=
 github.com/mozilla/tls-observatory v0.0.0-20210609171429-7bc42856d2e5/go.mod h1:FUqVoUPHSEdDR0MnFM3Dh8AU0pZHLXUD127SAJGER/s=
+github.com/mroth/weightedrand v0.4.1 h1:rHcbUBopmi/3x4nnrvwGJBhX9d0vk+KgoLUZeDP6YyI=
+github.com/mroth/weightedrand v0.4.1/go.mod h1:3p2SIcC8al1YMzGhAIoXD+r9olo/g/cdJgAD905gyNE=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-proto-validators v0.0.0-20180403085117-0950a7990007/go.mod h1:m2XC9Qq0AlmmVksL6FktJCdTYyLk7V3fKyp0sl1yWQo=
diff --git a/test/e2e/generator/generate.go b/test/e2e/generator/generate.go
index 12997eb81..1c48157a1 100644
--- a/test/e2e/generator/generate.go
+++ b/test/e2e/generator/generate.go
@@ -26,7 +26,13 @@ var (
 	}
 
 	// The following specify randomly chosen values for testnet nodes.
-	nodeDatabases        = uniformChoice{"goleveldb", "cleveldb", "rocksdb", "boltdb", "badgerdb"}
+	nodeDatabases = weightedChoice{
+		"goleveldb": 35,
+		"badgerdb":  35,
+		"boltdb":    15,
+		"rocksdb":   10,
+		"cleveldb":  5,
+	}
 	nodeABCIProtocols    = uniformChoice{"unix", "tcp", "builtin", "grpc"}
 	nodePrivvalProtocols = uniformChoice{"file", "unix", "tcp", "grpc"}
 	// FIXME: v2 disabled due to flake
@@ -270,7 +276,7 @@ func generateNode(
 	node := e2e.ManifestNode{
 		Mode:             string(mode),
 		StartAt:          startAt,
-		Database:         nodeDatabases.Choose(r).(string),
+		Database:         nodeDatabases.Choose(r),
 		ABCIProtocol:     nodeABCIProtocols.Choose(r).(string),
 		PrivvalProtocol:  nodePrivvalProtocols.Choose(r).(string),
 		BlockSync:        nodeBlockSyncs.Choose(r).(string),
@@ -321,7 +327,7 @@ func generateLightNode(r *rand.Rand, startAt int64, providers []string) *e2e.Man
 	return &e2e.ManifestNode{
 		Mode:            string(e2e.ModeLight),
 		StartAt:         startAt,
-		Database:        nodeDatabases.Choose(r).(string),
+		Database:        nodeDatabases.Choose(r),
 		ABCIProtocol:    "builtin",
 		PersistInterval: ptrUint64(0),
 		PersistentPeers: providers,
diff --git a/test/e2e/generator/random.go b/test/e2e/generator/random.go
index f21502118..d6c84d46c 100644
--- a/test/e2e/generator/random.go
+++ b/test/e2e/generator/random.go
@@ -3,6 +3,8 @@ package main
 import (
 	"math/rand"
 	"sort"
+
+	"github.com/mroth/weightedrand"
 )
 
 // combinations takes input in the form of a map of item lists, and returns a
@@ -83,3 +85,19 @@ func (usc uniformSetChoice) Choose(r *rand.Rand) []string {
 	}
 	return choices
 }
+
+type weightedChoice map[string]uint
+
+func (wc weightedChoice) Choose(r *rand.Rand) string {
+	choices := make([]weightedrand.Choice, 0, len(wc))
+	for k, v := range wc {
+		choices = append(choices, weightedrand.NewChoice(k, v))
+	}
+
+	chooser, err := weightedrand.NewChooser(choices...)
+	if err != nil {
+		panic(err)
+	}
+
+	return chooser.PickSource(r).(string)
+}