Merge pull request cockroachdb#112394 from RaduBerinde/backport23.1-1…

…10687

release-23.1:  roachtest: add CompatibleClouds and Suites to TestSpec

pkg/BUILD.bazel

@@ -142,6 +142,7 @@ ALL_TESTS = [
     "//pkg/cmd/roachprod-microbench:roachprod-microbench_test",
     "//pkg/cmd/roachtest/clusterstats:clusterstats_test",
     "//pkg/cmd/roachtest/option:option_test",
+    "//pkg/cmd/roachtest/registry:registry_test",
     "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test",
     "//pkg/cmd/roachtest/roachtestutil:roachtestutil_test",
     "//pkg/cmd/roachtest/tests:tests_test",
@@ -1074,6 +1075,7 @@ GO_TARGETS = [
     "//pkg/cmd/roachtest/option:option",
     "//pkg/cmd/roachtest/option:option_test",
     "//pkg/cmd/roachtest/registry:registry",
+    "//pkg/cmd/roachtest/registry:registry_test",
     "//pkg/cmd/roachtest/roachtestutil/clusterupgrade:clusterupgrade",
     "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion",
     "//pkg/cmd/roachtest/roachtestutil/mixedversion:mixedversion_test",

pkg/cmd/roachtest/main_test.go

@@ -37,10 +37,12 @@ func makeRegistry(names ...string) testRegistryImpl {
 
 	for _, name := range names {
 		r.Add(registry.TestSpec{
-			Name:    name,
-			Owner:   OwnerUnitTest,
-			Run:     dummyRun,
-			Cluster: spec.MakeClusterSpec(spec.GCE, "", 0),
+			Name:             name,
+			Owner:            OwnerUnitTest,
+			Run:              dummyRun,
+			Cluster:          spec.MakeClusterSpec(spec.GCE, "", 0),
+			CompatibleClouds: registry.AllExceptAWS,
+			Suites:           registry.Suites(registry.Nightly),
 		})
 	}
 

pkg/cmd/roachtest/registry/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "registry",
@@ -19,3 +19,14 @@ go_library(
         "@com_github_prometheus_client_golang//prometheus/promauto",
     ],
 )
+
+go_test(
+    name = "registry_test",
+    srcs = ["test_spec_test.go"],
+    args = ["-test.timeout=295s"],
+    embed = [":registry"],
+    deps = [
+        "//pkg/cmd/roachtest/spec",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/cmd/roachtest/registry/test_spec.go

@@ -54,6 +54,15 @@ type TestSpec struct {
 	// N.B. performance tests may have different requirements than correctness tests, e.g., machine type/architecture.
 	// Thus, they must be opted into explicitly via this field.
 	Benchmark bool
+
+	// CompatibleClouds is the set of clouds this test can run on (e.g. AllClouds,
+	// OnlyGCE, etc). Must be set.
+	CompatibleClouds CloudSet
+
+	// Suites is the set of suites this test is part of (e.g. Nightly, Weekly,
+	// etc). Must be set, even if empty (see ManualOnly).
+	Suites SuiteSet
+
 	// Tags is a set of tags associated with the test that allow grouping
 	// tests. If no tags are specified, the set ["default"] is automatically
 	// given.
@@ -196,6 +205,221 @@ const (
 	MetamorphicLeases
 )
 
+var allClouds = []string{spec.Local, spec.GCE, spec.AWS, spec.Azure}
+
+// CloudSet represents a set of clouds.
+//
+// Instances of CloudSet are immutable. The uninitialized (zero) value is not
+// valid.
+type CloudSet struct {
+	// m contains only values from allClouds.
+	m map[string]struct{}
+}
+
+// AllClouds contains all clouds.
+var AllClouds = Clouds(allClouds...)
+
+// AllExceptLocal contains all clouds except Local.
+var AllExceptLocal = AllClouds.NoLocal()
+
+// AllExceptAWS contains all clouds except AWS.
+var AllExceptAWS = AllClouds.NoAWS()
+
+// OnlyGCE contains only the GCE cloud.
+var OnlyGCE = Clouds(spec.GCE)
+
+// OnlyLocal contains only the GCE cloud.
+var OnlyLocal = Clouds(spec.Local)
+
+// Clouds creates a CloudSet for the given clouds. Cloud names must be one of:
+// spec.Local, spec.GCE, spec.AWS, spec.Azure.
+func Clouds(clouds ...string) CloudSet {
+	assertValidValues(allClouds, clouds...)
+	return CloudSet{m: addToSet(nil, clouds...)}
+}
+
+// NoLocal removes the Local cloud and returns the new set.
+func (cs CloudSet) NoLocal() CloudSet {
+	return CloudSet{m: removeFromSet(cs.m, spec.Local)}
+}
+
+// NoAWS removes the AWS cloud and returns the new set.
+func (cs CloudSet) NoAWS() CloudSet {
+	return CloudSet{m: removeFromSet(cs.m, spec.AWS)}
+}
+
+// NoAzure removes the Azure cloud and returns the new set.
+func (cs CloudSet) NoAzure() CloudSet {
+	return CloudSet{m: removeFromSet(cs.m, spec.Azure)}
+}
+
+// Contains returns true if the set contains the given cloud.
+func (cs CloudSet) Contains(cloud string) bool {
+	cs.AssertInitialized()
+	_, ok := cs.m[cloud]
+	return ok
+}
+
+func (cs CloudSet) String() string {
+	cs.AssertInitialized()
+	return setToString(allClouds, cs.m)
+}
+
+// AssertInitialized panics if the CloudSet is the zero value.
+func (cs CloudSet) AssertInitialized() {
+	if cs.m == nil {
+		panic("CloudSet not initialized")
+	}
+}
+
+// Suite names.
+const (
+	Nightly               = "nightly"
+	Weekly                = "weekly"
+	ReleaseQualification  = "release_qualification"
+	ORM                   = "orm"
+	Driver                = "driver"
+	Tool                  = "tool"
+	Smoketest             = "smoketest"
+	Quick                 = "quick"
+	Fixtures              = "fixtures"
+	Pebble                = "pebble"
+	PebbleNightlyWrite    = "pebble_nightly_write"
+	PebbleNightlyYCSB     = "pebble_nightly_ycsb"
+	PebbleNightlyYCSBRace = "pebble_nightly_ycsb_race"
+	Roachtest             = "roachtest"
+)
+
+var allSuites = []string{
+	Nightly, Weekly, ReleaseQualification, ORM, Driver, Tool, Smoketest, Quick, Fixtures,
+	Pebble, PebbleNightlyWrite, PebbleNightlyYCSB, PebbleNightlyYCSBRace, Roachtest,
+}
+
+// SuiteSet represents a set of suites.
+//
+// Instances of SuiteSet are immutable. The uninitialized (zero) value is not
+// valid.
+type SuiteSet struct {
+	// m contains only values from allSuites.
+	m map[string]struct{}
+}
+
+// ManualOnly is used for tests that are not part of any suite; these tests are
+// only run manually.
+var ManualOnly = Suites()
+
+// Suites creates a SuiteSet with the given suites. Only the constants above are
+// valid values.
+func Suites(suites ...string) SuiteSet {
+	assertValidValues(allSuites, suites...)
+	return SuiteSet{m: addToSet(nil, suites...)}
+}
+
+// Contains returns true if the set contains the given suite.
+func (ss SuiteSet) Contains(suite string) bool {
+	ss.AssertInitialized()
+	_, ok := ss.m[suite]
+	return ok
+}
+
+func (ss SuiteSet) String() string {
+	return setToString(allSuites, ss.m)
+}
+
+// AssertInitialized panics if the SuiteSet is the zero value.
+func (ss SuiteSet) AssertInitialized() {
+	if ss.m == nil {
+		panic("SuiteSet not initialized")
+	}
+}
+
+// CrossCheckTags verifies that the CompatibleClouds and Suites values match those of the tags.
+func (t *TestSpec) CrossCheckTags() {
+	// Check that CompatibleClouds + Suites would make the same determination as
+	// tags for what test to run for Nightly/Weekly and GCE/AWS.
+	var expected, actual struct {
+		nightlyGCE bool
+		weeklyGCE  bool
+		nightlyAWS bool
+		weeklyAWS  bool
+	}
+
+	expected.nightlyGCE = matchesAll(t.Tags, []string{"default"}) || matchesAll(t.Tags, []string{"aws"})
+	expected.weeklyGCE = matchesAll(t.Tags, []string{"weekly"})
+	expected.nightlyAWS = matchesAll(t.Tags, []string{"aws"})
+	expected.weeklyAWS = matchesAll(t.Tags, []string{"aws-weekly"})
+
+	actual.nightlyGCE = t.Suites.Contains(Nightly) && t.CompatibleClouds.Contains(spec.GCE)
+	actual.weeklyGCE = t.Suites.Contains(Weekly) && t.CompatibleClouds.Contains(spec.GCE)
+	actual.nightlyAWS = t.Suites.Contains(Nightly) && t.CompatibleClouds.Contains(spec.AWS)
+	actual.weeklyAWS = t.Suites.Contains(Weekly) && t.CompatibleClouds.Contains(spec.AWS)
+
+	if actual != expected {
+		panic(fmt.Sprintf("CompatibleClouds/Suites inconsistent with Tags\nexpected: %#v\nactual:   %#v\nclouds: %s  suites:%s  tags:%v\n", expected, actual, t.CompatibleClouds, t.Suites, t.Tags))
+	}
+
+	otherSuiteTags := fmt.Sprintf("%v", removeFromSet(t.Tags, "default", "weekly", "aws-weekly", "aws", "owner-"+string(t.Owner)))
+	otherSuites := fmt.Sprintf("%v", removeFromSet(t.Suites.m, Weekly, Nightly))
+	if otherSuiteTags != otherSuites && !(otherSuiteTags == "map[manual:{}]" && otherSuites == "map[]") {
+		panic(fmt.Sprintf("Suites inconsistent with Tags\nexpected: %v\nactual:   %v\nsuites:%s  tags:%v\n", otherSuiteTags, otherSuites, t.Suites, t.Tags))
+	}
+}
+
+// assertValidValues asserts that the given values exist in the validValues slice.
+func assertValidValues(validValues []string, values ...string) {
+	for _, v := range values {
+		found := false
+		for _, valid := range validValues {
+			if valid == v {
+				found = true
+				break
+			}
+		}
+		if !found {
+			panic(fmt.Sprintf("invalid value %q (valid values: %v)", v, validValues))
+		}
+	}
+}
+
+// addToSet returns a new set that is the initial set with the given values added.
+func addToSet(initial map[string]struct{}, values ...string) map[string]struct{} {
+	m := make(map[string]struct{})
+	for k := range initial {
+		m[k] = struct{}{}
+	}
+	for _, v := range values {
+		m[v] = struct{}{}
+	}
+	return m
+}
+
+// removeFromSet returns a new set that is the initial set with the given values removed.
+func removeFromSet(initial map[string]struct{}, values ...string) map[string]struct{} {
+	m := make(map[string]struct{})
+	for k := range initial {
+		m[k] = struct{}{}
+	}
+	for _, v := range values {
+		delete(m, v)
+	}
+	return m
+}
+
+// setToString returns the elements of a set, in the relative order in which they appear
+// in validValues. Returns "<none>" if the set is empty.
+func setToString(validValues []string, m map[string]struct{}) string {
+	var elems []string
+	for _, v := range validValues {
+		if _, ok := m[v]; ok {
+			elems = append(elems, v)
+		}
+	}
+	if len(elems) == 0 {
+		return "<none>"
+	}
+	return strings.Join(elems, ",")
+}
+
 func matchesAll(testTags map[string]struct{}, filterTags []string) bool {
 	for _, tag := range filterTags {
 		negate := false

pkg/cmd/roachtest/registry/test_spec_test.go

@@ -0,0 +1,45 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package registry
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCloudSet(t *testing.T) {
+	expect := func(c CloudSet, exp string) {
+		require.Equal(t, exp, c.String())
+	}
+	expect(AllClouds, "local,gce,aws,azure")
+	expect(AllExceptAWS, "local,gce,azure")
+	expect(AllExceptLocal, "gce,aws,azure")
+	expect(AllExceptLocal.NoAWS(), "gce,azure")
+	expect(AllClouds.NoAWS().NoAzure(), "local,gce")
+
+	require.True(t, AllExceptAWS.Contains(spec.GCE))
+	require.True(t, AllExceptAWS.Contains(spec.Local))
+	require.False(t, AllExceptAWS.Contains(spec.AWS))
+}
+
+func TestSuiteSet(t *testing.T) {
+	expect := func(c SuiteSet, exp string) {
+		require.Equal(t, exp, c.String())
+	}
+	s := Suites(Nightly, Weekly)
+	expect(s, "nightly,weekly")
+	require.True(t, s.Contains(Nightly))
+	require.True(t, s.Contains(Weekly))
+	require.False(t, s.Contains(ReleaseQualification))
+	expect(ManualOnly, "<none>")
+}

pkg/cmd/roachtest/test_registry.go

@@ -75,6 +75,7 @@ func (r *testRegistryImpl) Add(spec registry.TestSpec) {
 		fmt.Fprintf(os.Stderr, "%+v\n", err)
 		os.Exit(1)
 	}
+	spec.CrossCheckTags()
 	r.m[spec.Name] = &spec
 }
 
@@ -102,6 +103,9 @@ func (r *testRegistryImpl) prepareSpec(spec *registry.TestSpec) error {
 		return fmt.Errorf("%s: Name must match this regexp: %s", spec.Name, testNameRE)
 	}
 
+	spec.CompatibleClouds.AssertInitialized()
+	spec.Suites.AssertInitialized()
+
 	if spec.Run == nil {
 		return fmt.Errorf("%s: must specify Run", spec.Name)
 	}