From bbb4c25c35bf0dde205d9edc8d9317bcb88c273e Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@cockroachlabs.com>
Date: Mon, 9 Sep 2019 13:05:45 -0400
Subject: [PATCH] cli: Add some tests for the locality setting features of
 cockroach demo

Added some tests for locality settings in cockroach demo.

Release justification: Low risk tests for an existing feature.

Release note: None
---
 pkg/cli/cli_test.go                           | 40 +++++++++++++++++++
 .../test_demo_locality_error.tcl              | 19 +++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 pkg/cli/interactive_tests/test_demo_locality_error.tcl

diff --git a/pkg/cli/cli_test.go b/pkg/cli/cli_test.go
index ab631dba54e9..792b1e627f91 100644
--- a/pkg/cli/cli_test.go
+++ b/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)
 	}
diff --git a/pkg/cli/interactive_tests/test_demo_locality_error.tcl b/pkg/cli/interactive_tests/test_demo_locality_error.tcl
new file mode 100644
index 000000000000..4cedeb632039
--- /dev/null
+++ b/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