Merge cockroachdb#40602

40602: cli: Add some tests for the locality setting features of cockroach demo r=knz a=rohany

Added some tests for locality settings in cockroach demo.

Release justification: Low risk tests for an existing feature.

Release note: None

Co-authored-by: Rohan Yadav <rohany@cockroachlabs.com>

pkg/cli/cli_test.go

@@ -470,6 +470,45 @@ func Example_logging() {
 	// 1
 }
 
+func Example_demo_locality() {
+	c := newCLITest(cliTestParams{noServer: true})
+	defer c.cleanup()
+
+	testData := [][]string{
+		{`demo`, `--nodes`, `3`, `-e`, `select node_id, locality from crdb_internal.gossip_nodes order by node_id`},
+		{`demo`, `--nodes`, `9`, `-e`, `select node_id, locality from crdb_internal.gossip_nodes order by node_id`},
+		{`demo`, `--nodes`, `3`, `--demo-locality=region=us-east1:region=us-east2:region=us-east3`,
+			`-e`, `select node_id, locality from crdb_internal.gossip_nodes order by node_id`},
+	}
+	setCLIDefaultsForTests()
+	for _, cmd := range testData {
+		c.RunWithArgs(cmd)
+	}
+
+	// Output:
+	// demo --nodes 3 -e select node_id, locality from crdb_internal.gossip_nodes order by node_id
+	// node_id	locality
+	// 1	region=us-east1,az=b
+	// 2	region=us-east1,az=c
+	// 3	region=us-east1,az=d
+	// demo --nodes 9 -e select node_id, locality from crdb_internal.gossip_nodes order by node_id
+	// node_id	locality
+	// 1	region=us-east1,az=b
+	// 2	region=us-east1,az=c
+	// 3	region=us-east1,az=d
+	// 4	region=us-west1,az=a
+	// 5	region=us-west1,az=b
+	// 6	region=us-west1,az=c
+	// 7	region=europe-west1,az=b
+	// 8	region=europe-west1,az=c
+	// 9	region=europe-west1,az=d
+	// demo --nodes 3 --demo-locality=region=us-east1:region=us-east2:region=us-east3 -e select node_id, locality from crdb_internal.gossip_nodes order by node_id
+	// node_id	locality
+	// 1	region=us-east1
+	// 2	region=us-east2
+	// 3	region=us-east3
+}
+
 func Example_demo() {
 	c := newCLITest(cliTestParams{noServer: true})
 	defer c.cleanup()
@@ -485,6 +524,7 @@ func Example_demo() {
 		{`demo`, `startrek`, `-e`, `show databases`},
 		{`demo`, `startrek`, `-e`, `show databases`, `--format=table`},
 	}
+	setCLIDefaultsForTests()
 	for _, cmd := range testData {
 		c.RunWithArgs(cmd)
 	}

pkg/cli/interactive_tests/test_demo_locality_error.tcl

@@ -0,0 +1,19 @@
+#! /usr/bin/env expect -f
+
+source [file join [file dirname $argv0] common.tcl]
+
+start_test "Expect error with incorrect demo locality settings"
+
+# Test failure with less localities than expected.
+spawn $argv demo --nodes 3 --demo-locality=region=us-east1:region=us-east2
+# wait for the CLI to start up
+eexpect "Error: number of localities specified must equal number of nodes"
+eexpect eof
+
+# Test failure with more localities than expected.
+spawn $argv demo --nodes 3 --demo-locality=region=us-east1:region=us-east2:region=us-east3:region=us-east4
+# wait for the CLI to start up
+eexpect "Error: number of localities specified must equal number of nodes"
+eexpect eof
+
+end_test