Merge #156336

156336: testutils: introduce datadriven test helpers r=stevendanna a=pav-kv

This PR introduces 3 trivial helpers for parsing arguments in `datadriven` tests. These cover the vast majority of use-cases and cut the parsing boilerplate in half. See individual commits, separated by package (only sub-packages of `kvserver` at this point).

Examples:

```golang
var (
	rangeID = dd.ScanArg[roachpb.RangeID](t, d, "id") // required
	flag    = dd.ScanArgOr(t, d, "flag", false)       // optional flag
	param   = dd.ScanArgOr[uint32](t, d, "param", 10) // optional, with a default
	ids     = dd.ScanArgOr[[]roachpb.NodeID](t, d, "nodes", nil)  // slice
)

// Optional argument with explicit presence handling.
if str, ok := dd.ScanArgOpt[string](t, d, "key"); ok {
	// do something with str, e.g. parse further
}
```

This PR also simplifies a few places in tests where manual parsing was unnecessary because `datadriven` supports the corresponding types natively (e.g. `time.Duration` and slices).

In many cases, it also removes redundant type conversions, since the right type can be used as the generic parameter from the beginning.

Epic: none
Release note: none

Co-authored-by: Pavel Kalinnikov <pavel@cockroachlabs.com>

Author: craig[bot]

pkg/BUILD.bazel

@@ -2472,6 +2472,7 @@ GO_TARGETS = [
     "//pkg/testutils/bazelcodecover:bazelcodecover",
     "//pkg/testutils/colcontainerutils:colcontainerutils",
     "//pkg/testutils/datapathutils:datapathutils",
+    "//pkg/testutils/dd:dd",
     "//pkg/testutils/diagutils:diagutils",
     "//pkg/testutils/distsqlutils:distsqlutils",
     "//pkg/testutils/docker/docker-fsnotify:docker-fsnotify",

pkg/kv/kvserver/BUILD.bazel

@@ -517,6 +517,7 @@ go_test(
         "//pkg/storage/storageconfig",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",
+        "//pkg/testutils/dd",
         "//pkg/testutils/echotest",
         "//pkg/testutils/gossiputil",
         "//pkg/testutils/jobutils",

pkg/kv/kvserver/allocator/mmaprototype/BUILD.bazel

@@ -47,6 +47,7 @@ go_test(
         "//pkg/roachpb",
         "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/testutils/datapathutils",
+        "//pkg/testutils/dd",
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_datadriven//:datadriven",

pkg/kv/kvserver/allocator/mmaprototype/allocator_state_test.go

@@ -8,10 +8,10 @@ package mmaprototype
 import (
 	"fmt"
 	"slices"
-	"strconv"
 	"strings"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/datadriven"
 	"github.com/stretchr/testify/require"
 )
@@ -26,16 +26,8 @@ func TestDiversityScoringMemo(t *testing.T) {
 	datadriven.RunTest(t, "testdata/diversity_scoring_memo",
 		func(t *testing.T, d *datadriven.TestData) string {
 			scanStores := func() []localityTiers {
-				var storesStr string
-				d.ScanArgs(t, "store-ids", &storesStr)
-				storesStrSlice := strings.Split(storesStr, ",")
-				var storeIDs []int
-				for _, s := range storesStrSlice {
-					storeID, err := strconv.Atoi(s)
-					require.NoError(t, err)
-					storeIDs = append(storeIDs, storeID)
-				}
-				var res []localityTiers
+				storeIDs := dd.ScanArg[[]int](t, d, "store-ids")
+				res := make([]localityTiers, 0, len(storeIDs))
 				for _, storeID := range storeIDs {
 					l, ok := storeLocalities[storeID]
 					require.True(t, ok)
@@ -59,18 +51,15 @@ func TestDiversityScoringMemo(t *testing.T) {
 				}
 			}
 			getStoreLocality := func(key string) localityTiers {
-				var storeID int
-				d.ScanArgs(t, key, &storeID)
+				storeID := dd.ScanArg[int](t, d, key)
 				l, ok := storeLocalities[storeID]
 				require.True(t, ok)
 				return l
 			}
 			switch d.Cmd {
 			case "store":
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
-				var lts string
-				d.ScanArgs(t, "locality-tiers", &lts)
+				storeID := dd.ScanArg[int](t, d, "store-id")
+				lts := dd.ScanArg[string](t, d, "locality-tiers")
 				locality := parseLocalityTiers(t, lts)
 				lt := ltInterner.intern(locality)
 				storeLocalities[storeID] = lt

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigtestutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/datadriven"
 	"github.com/stretchr/testify/require"
@@ -365,20 +366,17 @@ func TestClusterState(t *testing.T) {
 					return printNodeListMeta()
 
 				case "set-store-membership":
-					var storeID int
-					d.ScanArgs(t, "store-id", &storeID)
-					var storeMembershipString string
-					d.ScanArgs(t, "membership", &storeMembershipString)
+					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 					var storeMembershipVal storeMembership
-					switch storeMembershipString {
+					switch str := dd.ScanArg[string](t, d, "membership"); str {
 					case "member":
 						storeMembershipVal = storeMembershipMember
 					case "removing":
 						storeMembershipVal = storeMembershipRemoving
 					case "removed":
 						storeMembershipVal = storeMembershipRemoved
 					}
-					cs.setStoreMembership(roachpb.StoreID(storeID), storeMembershipVal)
+					cs.setStoreMembership(storeID, storeMembershipVal)
 
 					var buf strings.Builder
 					nonRemovedStores, removedStores := testingGetStoreList(t, cs)
@@ -389,18 +387,16 @@ func TestClusterState(t *testing.T) {
 					return buf.String()
 
 				case "update-failure-detection":
-					var nodeID int
-					var failureDetectionString string
-					d.ScanArgs(t, "node-id", &nodeID)
-					d.ScanArgs(t, "summary", &failureDetectionString)
+					nodeID := dd.ScanArg[roachpb.NodeID](t, d, "node-id")
+					failureDetectionString := dd.ScanArg[string](t, d, "summary")
 					var fd failureDetectionSummary
 					for i := fdOK; i < fdDead+1; i++ {
 						if i.String() == failureDetectionString {
 							fd = i
 							break
 						}
 					}
-					cs.updateFailureDetectionSummary(roachpb.NodeID(nodeID), fd)
+					cs.updateFailureDetectionSummary(nodeID, fd)
 					return printNodeListMeta()
 
 				case "store-load-msg":
@@ -414,10 +410,8 @@ func TestClusterState(t *testing.T) {
 					return ""
 
 				case "make-pending-changes":
-					var rid int
+					rangeID := dd.ScanArg[roachpb.RangeID](t, d, "range-id")
 					var changes []ReplicaChange
-					d.ScanArgs(t, "range-id", &rid)
-					rangeID := roachpb.RangeID(rid)
 					rState := cs.ranges[rangeID]
 
 					lines := strings.Split(d.Input, "\n")
@@ -467,19 +461,17 @@ func TestClusterState(t *testing.T) {
 					return printPendingChangesTest(testingGetPendingChanges(t, cs))
 
 				case "reject-pending-changes":
-					var changeIDsInt []int
-					d.ScanArgs(t, "change-ids", &changeIDsInt)
+					changeIDsInt := dd.ScanArg[[]ChangeID](t, d, "change-ids")
 					for _, id := range changeIDsInt {
-						cs.undoPendingChange(ChangeID(id), true)
+						cs.undoPendingChange(id, true)
 					}
 					return printPendingChangesTest(testingGetPendingChanges(t, cs))
 
 				case "get-pending-changes":
 					return printPendingChangesTest(testingGetPendingChanges(t, cs))
 
 				case "tick":
-					var seconds int
-					d.ScanArgs(t, "seconds", &seconds)
+					seconds := dd.ScanArg[int](t, d, "seconds")
 					ts.Advance(time.Second * time.Duration(seconds))
 					return fmt.Sprintf("t=%v", ts.Now().Sub(testingBaseTime))
 

pkg/kv/kvserver/allocator/mmaprototype/constraint_matcher_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/datadriven"
 	"github.com/stretchr/testify/require"
@@ -99,23 +100,21 @@ func TestConstraintMatcher(t *testing.T) {
 				return b.String()
 
 			case "remove-store":
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
-				cm.removeStore(roachpb.StoreID(storeID))
+				storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
+				cm.removeStore(storeID)
 				var b strings.Builder
 				printMatcher(&b)
 				return b.String()
 
 			case "store-matches":
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
+				storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 				lines := strings.Split(d.Input, "\n")
 				require.Greater(t, 2, len(lines))
 				var cc []roachpb.Constraint
 				if len(lines) == 1 {
 					cc = parseConstraints(t, strings.Fields(strings.TrimSpace(lines[0])))
 				}
-				matches := cm.storeMatches(roachpb.StoreID(storeID), interner.internConstraintsConj(cc))
+				matches := cm.storeMatches(storeID, interner.internConstraintsConj(cc))
 				return fmt.Sprintf("%t", matches)
 
 			case "match-stores":

pkg/kv/kvserver/allocator/mmaprototype/constraint_test.go

@@ -14,6 +14,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
@@ -66,12 +67,10 @@ func parseConstraintsConj(t *testing.T, fields []string) roachpb.ConstraintsConj
 }
 
 func parseSpanConfig(t *testing.T, d *datadriven.TestData) roachpb.SpanConfig {
-	var numReplicas, numVoters int
-	var conf roachpb.SpanConfig
-	d.ScanArgs(t, "num-replicas", &numReplicas)
-	conf.NumReplicas = int32(numReplicas)
-	d.ScanArgs(t, "num-voters", &numVoters)
-	conf.NumVoters = int32(numVoters)
+	conf := roachpb.SpanConfig{
+		NumReplicas: dd.ScanArg[int32](t, d, "num-replicas"),
+		NumVoters:   dd.ScanArg[int32](t, d, "num-voters"),
+	}
 	for _, line := range strings.Split(d.Input, "\n") {
 		parts := strings.Fields(line)
 		if len(parts) == 0 {
@@ -169,8 +168,7 @@ func TestStoreIDPostingList(t *testing.T) {
 		func(t *testing.T, d *datadriven.TestData) string {
 			switch d.Cmd {
 			case "pl":
-				var name string
-				d.ScanArgs(t, "name", &name)
+				name := dd.ScanArg[string](t, d, "name")
 				var storeIDs []roachpb.StoreID
 				for _, line := range strings.Split(d.Input, "\n") {
 					parts := strings.Fields(line)
@@ -190,9 +188,8 @@ func TestStoreIDPostingList(t *testing.T) {
 				return b.String()
 
 			case "intersect", "union", "is-equal":
-				var x, y string
-				d.ScanArgs(t, "x", &x)
-				d.ScanArgs(t, "y", &y)
+				x := dd.ScanArg[string](t, d, "x")
+				y := dd.ScanArg[string](t, d, "y")
 				plX := pls[x]
 				if d.Cmd == "is-equal" {
 					return fmt.Sprintf("%t", plX.isEqual(pls[y]))
@@ -212,19 +209,17 @@ func TestStoreIDPostingList(t *testing.T) {
 				}
 
 			case "insert", "contains", "remove":
-				var name string
-				d.ScanArgs(t, "name", &name)
+				name := dd.ScanArg[string](t, d, "name")
 				pl := pls[name]
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
+				storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 				if d.Cmd == "contains" {
-					return fmt.Sprintf("%t", pl.contains(roachpb.StoreID(storeID)))
+					return fmt.Sprintf("%t", pl.contains(storeID))
 				} else {
 					var rv bool
 					if d.Cmd == "insert" {
-						rv = pl.insert(roachpb.StoreID(storeID))
+						rv = pl.insert(storeID)
 					} else {
-						rv = pl.remove(roachpb.StoreID(storeID))
+						rv = pl.remove(storeID)
 					}
 					if forceAllocation {
 						pl = pl[:len(pl):len(pl)]
@@ -237,9 +232,7 @@ func TestStoreIDPostingList(t *testing.T) {
 				}
 
 			case "hash":
-				var name string
-				d.ScanArgs(t, "name", &name)
-				pl := pls[name]
+				pl := pls[dd.ScanArg[string](t, d, "name")]
 				return fmt.Sprintf("%d", pl.hash())
 
 			default:
@@ -375,8 +368,7 @@ func TestRangeAnalyzedConstraints(t *testing.T) {
 				return ""
 
 			case "span-config":
-				var name string
-				d.ScanArgs(t, "name", &name)
+				name := dd.ScanArg[string](t, d, "name")
 				conf := parseSpanConfig(t, d)
 				var b strings.Builder
 				nConf, err := makeNormalizedSpanConfig(&conf, interner)
@@ -388,10 +380,8 @@ func TestRangeAnalyzedConstraints(t *testing.T) {
 				return b.String()
 
 			case "analyze-constraints":
-				var configName string
-				d.ScanArgs(t, "config-name", &configName)
-				var leaseholder int
-				d.ScanArgs(t, "leaseholder", &leaseholder)
+				configName := dd.ScanArg[string](t, d, "config-name")
+				leaseholder := dd.ScanArg[roachpb.StoreID](t, d, "leaseholder")
 				nConf := configs[configName]
 				rac := rangeAnalyzedConstraintsPool.Get().(*rangeAnalyzedConstraints)
 				buf := rac.stateForInit()
@@ -419,7 +409,7 @@ func TestRangeAnalyzedConstraints(t *testing.T) {
 					buf.tryAddingStore(roachpb.StoreID(storeID), typ,
 						ltInterner.intern(stores[roachpb.StoreID(storeID)].locality()))
 				}
-				rac.finishInit(nConf, cm, roachpb.StoreID(leaseholder))
+				rac.finishInit(nConf, cm, leaseholder)
 				var b strings.Builder
 				printRangeAnalyzedConstraints(&b, rac, ltInterner)
 				// If there is a previous rangeAnalyzedConstraints, release it before

pkg/kv/kvserver/allocator/mmaprototype/load_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/datadriven"
 	"github.com/stretchr/testify/require"
 )
@@ -75,44 +76,37 @@ func TestMeansMemo(t *testing.T) {
 				return ""
 
 			case "store-load":
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
-				sal, ok := storeMap[roachpb.StoreID(storeID)]
+				storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
+				sal, ok := storeMap[storeID]
 				require.True(t, ok)
 				var cpuLoad, wbLoad, bsLoad int64
 				d.ScanArgs(t, "load", &cpuLoad, &wbLoad, &bsLoad)
 				var cpuCapacity, wbCapacity, bsCapacity int64
 				d.ScanArgs(t, "capacity", &cpuCapacity, &wbCapacity, &bsCapacity)
-				var leaseCountLoad int64
-				d.ScanArgs(t, "secondary-load", &leaseCountLoad)
+				leaseCountLoad := dd.ScanArg[LoadValue](t, d, "secondary-load")
 				sLoad := &storeLoad{
 					reportedLoad: LoadVector{LoadValue(cpuLoad), LoadValue(wbLoad), LoadValue(bsLoad)},
 					capacity: LoadVector{
 						LoadValue(cpuCapacity), LoadValue(wbCapacity), LoadValue(bsCapacity)},
-					reportedSecondaryLoad: SecondaryLoadVector{LoadValue(leaseCountLoad)},
+					reportedSecondaryLoad: SecondaryLoadVector{leaseCountLoad},
 				}
 				for i := range sLoad.capacity {
 					if sLoad.capacity[i] < 0 {
 						sLoad.capacity[i] = UnknownCapacity
 					}
 				}
-				loadProvider.sloads[roachpb.StoreID(storeID)] = storeLoadAndNodeID{
+				loadProvider.sloads[storeID] = storeLoadAndNodeID{
 					nodeID:    sal.NodeID,
 					storeLoad: sLoad,
 				}
 
 				return ""
 
 			case "node-load":
-				var nodeID int
-				d.ScanArgs(t, "node-id", &nodeID)
-				var cpuLoad, cpuCapacity int64
-				d.ScanArgs(t, "cpu-load", &cpuLoad)
-				d.ScanArgs(t, "cpu-capacity", &cpuCapacity)
 				nLoad := &NodeLoad{
-					NodeID:      roachpb.NodeID(nodeID),
-					ReportedCPU: LoadValue(cpuLoad),
-					CapacityCPU: LoadValue(cpuCapacity),
+					NodeID:      dd.ScanArg[roachpb.NodeID](t, d, "node-id"),
+					ReportedCPU: dd.ScanArg[LoadValue](t, d, "cpu-load"),
+					CapacityCPU: dd.ScanArg[LoadValue](t, d, "cpu-capacity"),
 				}
 				loadProvider.nloads[nLoad.NodeID] = nLoad
 				return ""
@@ -157,12 +151,10 @@ func TestMeansMemo(t *testing.T) {
 				return b.String()
 
 			case "get-store-summary":
-				var storeID int
-				d.ScanArgs(t, "store-id", &storeID)
-				var loadSeqNum uint64
-				d.ScanArgs(t, "load-seq-num", &loadSeqNum)
+				storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
+				loadSeqNum := dd.ScanArg[uint64](t, d, "load-seq-num")
 				loadProvider.returnedLoadSeqNum = loadSeqNum
-				_ = mm.getStoreLoadSummary(context.Background(), mss, roachpb.StoreID(storeID), loadSeqNum)
+				_ = mm.getStoreLoadSummary(context.Background(), mss, storeID, loadSeqNum)
 				rv := loadProvider.b.String()
 				loadProvider.b.Reset()
 				return rv