From 558bee73bf0634385a862b19708de4ed19316fbb Mon Sep 17 00:00:00 2001
From: Jeff <swenson@cockroachlabs.com>
Date: Thu, 13 Apr 2023 16:49:33 -0400
Subject: [PATCH] testserver: add support for the --locality flag

This commit adds the `testserver.LocalityFlagsOpt` option. The flag
makes it possible to test multi-region primitives with the test server.
---
 testserver/testserver.go      | 20 ++++++++++++++++++++
 testserver/testserver_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/testserver/testserver.go b/testserver/testserver.go
index 3a06175..12c2e42 100644
--- a/testserver/testserver.go
+++ b/testserver/testserver.go
@@ -238,6 +238,7 @@ type testServerArgs struct {
 	initTimeoutSeconds          int
 	pollListenURLTimeoutSeconds int
 	envVars                     []string // to be passed to cmd.Env
+	localityFlags               []string
 }
 
 // CockroachBinaryPathOpt is a TestServer option that can be passed to
@@ -383,6 +384,18 @@ func PollListenURLTimeoutOpt(timeout int) TestServerOpt {
 	}
 }
 
+// LocalityFlagsOpt is used to specify the --locality flag for each node.
+//
+// Example Usage:
+//
+//	localities := LocalityFlagsOpt("region=us-west", "region=us-east", "region=us-central")
+//	server, err := NewTestServer(ThreeNodeOpt(), localities)
+func LocalityFlagsOpt(locality ...string) TestServerOpt {
+	return func(args *testServerArgs) {
+		args.localityFlags = locality
+	}
+}
+
 // EnvVarOpt is a list of environment variables to be passed to the start
 // command. Each entry in the slice should be in `key=value` format.
 func EnvVarOpt(vars []string) TestServerOpt {
@@ -430,6 +443,10 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 			serverArgs.numNodes, len(serverArgs.listenAddrPorts)))
 	}
 
+	if 0 < len(serverArgs.localityFlags) && len(serverArgs.localityFlags) != serverArgs.numNodes {
+		panic(fmt.Sprintf("got %d locality flags when %d are needed (one for each node)", len(serverArgs.localityFlags), serverArgs.numNodes))
+	}
+
 	var err error
 	if serverArgs.cockroachBinary != "" {
 		log.Printf("Using custom cockroach binary: %s", serverArgs.cockroachBinary)
@@ -574,6 +591,9 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 				"--external-io-dir=" + serverArgs.externalIODir,
 			}
 		}
+		if 0 < len(serverArgs.localityFlags) {
+			nodes[i].startCmdArgs = append(nodes[i].startCmdArgs, fmt.Sprintf("--locality=%s", serverArgs.localityFlags[i]))
+		}
 	}
 
 	// We only need initArgs if we're creating a testserver
diff --git a/testserver/testserver_test.go b/testserver/testserver_test.go
index 6b5b04b..ce799bd 100644
--- a/testserver/testserver_test.go
+++ b/testserver/testserver_test.go
@@ -784,3 +784,34 @@ func removeExistingLocalFile(localFile string) error {
 	}
 	return nil
 }
+
+func TestLocalityFlagsOpt(t *testing.T) {
+	ts, err := testserver.NewTestServer(
+		testserver.ThreeNodeOpt(),
+		testserver.LocalityFlagsOpt("region=us-east1", "region=us-central1", "region=us-west1"))
+	require.NoError(t, err)
+
+	for i := 0; i < 3; i++ {
+		ts.WaitForInitFinishForNode(i)
+	}
+
+	db, err := sql.Open("postgres", ts.PGURL().String())
+	require.NoError(t, err)
+
+	found := map[string]bool{}
+
+	rows, err := db.Query("SELECT region FROM [SHOW REGIONS]")
+	require.NoError(t, err)
+	defer rows.Close()
+	for rows.Next() {
+		var region string
+		require.NoError(t, rows.Scan(&region))
+		found[region] = true
+	}
+
+	require.Equal(t, map[string]bool{
+		"us-east1":    true,
+		"us-central1": true,
+		"us-west1":    true,
+	}, found)
+}