diff --git a/go.mod b/go.mod index 4ea2d2a38710..5d0cc0706d7e 100644 --- a/go.mod +++ b/go.mod @@ -94,6 +94,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v0.6.1 github.com/Azure/go-autorest/autorest/adal v0.9.15 github.com/BurntSushi/toml v1.2.1 + github.com/DATA-DOG/go-sqlmock v1.5.0 github.com/DataDog/datadog-api-client-go/v2 v2.15.0 github.com/DataExMachina-dev/side-eye-go v0.0.0-20240528211710-5eb9c7a69e1d github.com/IBM/sarama v1.42.1 diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index d91dd23f77db..ab6917f9bcdf 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -162,6 +162,7 @@ ALL_TESTS = [ "//pkg/cmd/roachtest/roachtestutil:roachtestutil_test", "//pkg/cmd/roachtest/spec:spec_test", "//pkg/cmd/roachtest/tests:tests_test", + "//pkg/cmd/roachtest/testselector:testselector_test", "//pkg/cmd/roachtest:roachtest_test", "//pkg/col/coldata:coldata_disallowed_imports_test", "//pkg/col/coldata:coldata_test", @@ -1222,6 +1223,7 @@ GO_TARGETS = [ "//pkg/cmd/roachtest/tests:tests", "//pkg/cmd/roachtest/tests:tests_test", "//pkg/cmd/roachtest/testselector:testselector", + "//pkg/cmd/roachtest/testselector:testselector_test", "//pkg/cmd/roachtest:roachtest", "//pkg/cmd/roachtest:roachtest_lib", "//pkg/cmd/roachtest:roachtest_test", diff --git a/pkg/cmd/roachtest/BUILD.bazel b/pkg/cmd/roachtest/BUILD.bazel index 14962aad3cbe..b9687d35c0a5 100644 --- a/pkg/cmd/roachtest/BUILD.bazel +++ b/pkg/cmd/roachtest/BUILD.bazel @@ -109,6 +109,7 @@ go_test( "//pkg/cmd/roachtest/roachtestflags", "//pkg/cmd/roachtest/spec", "//pkg/cmd/roachtest/test", + "//pkg/cmd/roachtest/testselector", "//pkg/internal/team", "//pkg/roachprod", "//pkg/roachprod/cloud", @@ -125,6 +126,7 @@ go_test( "//pkg/util/version", "@com_github_cockroachdb_datadriven//:datadriven", "@com_github_cockroachdb_errors//:errors", + "@com_github_data_dog_go_sqlmock//:go-sqlmock", "@com_github_kr_pretty//:pretty", "@com_github_prometheus_client_golang//prometheus", "@com_github_stretchr_testify//assert", diff --git a/pkg/cmd/roachtest/main_test.go b/pkg/cmd/roachtest/main_test.go index e7a62bf6a809..f25360077cc2 100644 --- a/pkg/cmd/roachtest/main_test.go +++ b/pkg/cmd/roachtest/main_test.go @@ -12,15 +12,22 @@ package main import ( "context" + gosql "database/sql" "fmt" + "os" + "regexp" "strings" "testing" + "github.com/DATA-DOG/go-sqlmock" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec" "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/test" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/testselector" "github.com/cockroachdb/cockroach/pkg/internal/team" + "github.com/stretchr/testify/require" ) func init() { @@ -98,3 +105,140 @@ func TestSampleSpecs(t *testing.T) { }) } } + +func Test_updateSpecForSelectiveTests(t *testing.T) { + ctx := context.Background() + var mock sqlmock.Sqlmock + var db *gosql.DB + var err error + _ = os.Setenv("SFUSER", "dummy_user") + _ = os.Setenv("SFPASSWORD", "dummy_password") + testselector.SqlConnectorFunc = func(_, _ string) (*gosql.DB, error) { + return db, err + } + t.Run("expect CategoriseTests to fail which causes fall back to run all tests", func(t *testing.T) { + // The failure of the CategoriseTests does not cause any failure. It falls back to run all the tests in the spec + db, mock, err = sqlmock.New() + require.Nil(t, err) + specs, _ := getTestSelectionMockData() + mock.ExpectPrepare(regexp.QuoteMeta(testselector.PreparedQuery)).WillReturnError(fmt.Errorf("failed to prepare")) + updateSpecForSelectiveTests(ctx, specs) + for _, s := range specs { + if !strings.Contains(s.Name, "skipped") { + require.Empty(t, s.Skip) + } else { + require.Equal(t, "test spec skip", s.Skip) + } + } + }) + t.Run("expect no failure", func(t *testing.T) { + db, mock, err = sqlmock.New() + require.Nil(t, err) + oldSuite := roachtestflags.Suite + roachtestflags.Suite = registry.Nightly + defer func() { + roachtestflags.Suite = oldSuite + }() + oldSuccessfulTestsSelectPct := roachtestflags.SuccessfulTestsSelectPct + roachtestflags.SuccessfulTestsSelectPct = 0.30 + defer func() { + roachtestflags.SuccessfulTestsSelectPct = oldSuccessfulTestsSelectPct + }() + specs, rows := getTestSelectionMockData() + mock.ExpectPrepare(regexp.QuoteMeta(testselector.PreparedQuery)) + mock.ExpectQuery(regexp.QuoteMeta(testselector.PreparedQuery)).WillReturnRows(rows) + specsLengthBefore := len(specs) + updateSpecForSelectiveTests(ctx, specs) + require.Equal(t, specsLengthBefore, len(specs)) + for _, s := range specs { + if strings.Contains(s.Name, "success_skip_selector") { + require.Equal(t, "test selector", s.Skip, s.Name) + require.Equal(t, "test skipped because it is stable and selective-tests is set.", s.SkipDetails, s.Name) + } else if strings.Contains(s.Name, "skipped") { + require.Equal(t, "test spec skip", s.Skip, s.Name) + require.Equal(t, "test spec skip test", s.SkipDetails, s.Name) + } else { + require.Empty(t, s.Skip, s.Name) + } + if strings.Contains(s.Name, "_preempted") { + require.True(t, s.IsLastFailurePreempt(), s.Name) + } else { + require.False(t, s.IsLastFailurePreempt(), s.Name) + } + } + }) +} + +// getTestSelectionMockData returns the mock data as: +// 1. List of test specs. +// 2. List of rows to be returned by snowflake. +// The data follows a convention: +// all tests which contains "_skipped" are skipped in the spec +// all tests which contains "_selected" are selected="yes" +// all tests which contains "_preempted" are last_failure_is_preempt="yes" +// all tests which contains "_skip_selector" are skipped based on the criteria +// all tests which contains "_missing" are missing in the test spec +// all tests which contains "_success" are selected="no" +// all tests which contains "new_test" are present in specs, but missing in snowflake rows +// all tests which contains "_randomized" are set in spec as Randomized=true +// Other than the above conventions, the name of the test is randomly put to represent random names for tests. +// the test names are also shuffled to ensure that the test is more close to real data. +func getTestSelectionMockData() ([]registry.TestSpec, *sqlmock.Rows) { + specs := []registry.TestSpec{ + {Name: "t2_selected_preempted"}, // selected by test selector for selected="yes" + {Name: "t_skipped_selected", Skip: "test spec skip", SkipDetails: "test spec skip test"}, + {Name: "t1_success"}, // selected by test selector based on the percentage criteria + {Name: "t_randomized", Randomized: true}, // not skipped by test selector as Randomized=true + {Name: "t_new_test_missing_in_sf_query_1"}, // selected even as this is missing in sf query result + {Name: "t2_success"}, // selected by test selector based on the percentage criteria + {Name: "t3_success"}, // selected by test selector based on the percentage criteria + {Name: "t_opt_out", TestSelectionOptOutSuites: registry.Suites(registry.Nightly)}, // selected as this test is opted out of Nightly + // opt out suite does not match, so, even if opt out is mentioned, the test is skipped + {Name: "t6_success_skip_selector", TestSelectionOptOutSuites: registry.Suites(registry.Weekly)}, + {Name: "t_skipped_new_test", Skip: "test spec skip", SkipDetails: "test spec skip test"}, // skipped in spec + {Name: "t_skipped_not_selected", Skip: "test spec skip", SkipDetails: "test spec skip test"}, // skipped in spec + {Name: "t5_success_skip_selector"}, // skipped by test selector based on the percentage criteria + {Name: "t1_selected"}, // selected by test selector for selected="yes" + {Name: "t8_success_skip_selector"}, // skipped by test selector based on the percentage criteria + {Name: "t_new_test_missing_in_sf_query"}, // selected even as this is missing in sf query result + {Name: "t10_success_skip_selector"}, // skipped by test selector based on the percentage criteria + } + for _, s := range specs { + s.Suites = registry.Suites(registry.Nightly, registry.Weekly) + } + // data are the rows returned by snowflake. This is a 2-dimensional list of strings represent the rows + // and columns that are returned. Each column has values for 4 rows in sequence: + // "name", "selected", "avg_duration", "last_failure_is_preempt" + data := [][]string{ + {"t1_selected", "yes", "1", "no"}, + {"t2_selected_preempted", "yes", "2", "yes"}, + {"t_skipped_selected", "yes", "3", "no"}, + // tests from here are selected under percent criteria + {"t1_success", "no", "4", "no"}, + {"t2_success", "no", "6", "no"}, + {"t3_success", "no", "7", "no"}, + // test is opted out from test selection + {"t_opt_out", "no", "8", "no"}, + {"t4_missing", "no", "9", "no"}, + // tests from here are skipped as these are beyond the percentage criteria + {"t5_success_skip_selector", "no", "10", "no"}, + {"t6_success_skip_selector", "no", "11", "no"}, + // the test is skipped by test selector, but is already skipped in spec + {"t_skipped_not_selected", "no", "12", "no"}, + {"t7_missing", "no", "13", "no"}, + {"t8_success_skip_selector", "no", "14", "no"}, + {"t9_missing", "no", "15", "no"}, + {"t10_success_skip_selector", "no", "16", "no"}, + // test is Randomized=true, so, will always be selected even being at the last in the list + {"t_randomized", "no", "5", "no"}, + } + // rows are the rows that are returned by snowflake. The list of string represents the columns of each row + rows := sqlmock.NewRows([]string{ + "name", "selected", "avg_duration", "last_failure_is_preempt", + }) + // rows are populated with the data + for _, ds := range data { + rows.FromCSVString(strings.Join(ds, ",")) + } + return specs, rows +} diff --git a/pkg/cmd/roachtest/testselector/BUILD.bazel b/pkg/cmd/roachtest/testselector/BUILD.bazel index 6ff71d235139..a6f90d799e76 100644 --- a/pkg/cmd/roachtest/testselector/BUILD.bazel +++ b/pkg/cmd/roachtest/testselector/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 = "testselector", @@ -11,3 +11,14 @@ go_library( "@com_github_snowflakedb_gosnowflake//:gosnowflake", ], ) + +go_test( + name = "testselector_test", + srcs = ["selector_test.go"], + embed = [":testselector"], + deps = [ + "//pkg/cmd/roachtest/spec", + "@com_github_data_dog_go_sqlmock//:go-sqlmock", + "@com_github_stretchr_testify//require", + ], +) diff --git a/pkg/cmd/roachtest/testselector/selector.go b/pkg/cmd/roachtest/testselector/selector.go index bc8210aef219..1a9fef994fd5 100644 --- a/pkg/cmd/roachtest/testselector/selector.go +++ b/pkg/cmd/roachtest/testselector/selector.go @@ -38,7 +38,10 @@ const ( ) //go:embed snowflake_query.sql -var preparedQuery string +var PreparedQuery string + +// SqlConnectorFunc is the function to get a sql connector +var SqlConnectorFunc = gosql.Open // supported suites var suites = map[string]string{ @@ -87,10 +90,11 @@ func CategoriseTests(ctx context.Context, req *SelectTestsReq) ([]*TestDetails, return nil, err } defer func() { _ = db.Close() }() - statement, err := db.Prepare(preparedQuery) + statement, err := db.Prepare(PreparedQuery) if err != nil { return nil, err } + defer func() { _ = statement.Close() }() // get the current branch from the teamcity environment currentBranch := os.Getenv("TC_BUILD_BRANCH") if currentBranch == "" { @@ -103,6 +107,7 @@ func CategoriseTests(ctx context.Context, req *SelectTestsReq) ([]*TestDetails, if err != nil { return nil, err } + defer func() { _ = rows.Close() }() // All the column headers colHeaders, err := rows.Columns() if err != nil { @@ -146,13 +151,26 @@ func getDuration(durationStr string) int64 { } // getConnect makes connection to snowflake and returns the connection. -func getConnect(ctx context.Context) (*gosql.DB, error) { - username, password, err := getSFCreds() +func getConnect(_ context.Context) (*gosql.DB, error) { + dsn, err := getDSN() + if err != nil { + return nil, err + } + db, err := SqlConnectorFunc("snowflake", dsn) if err != nil { return nil, err } + return db, nil +} + +// getDSN returns the dataSource name for snowflake driver +func getDSN() (string, error) { + username, password, err := getSFCreds() + if err != nil { + return "", err + } - dsn, err := sf.DSN(&sf.Config{ + return sf.DSN(&sf.Config{ Account: account, Database: database, Schema: schema, @@ -160,14 +178,6 @@ func getConnect(ctx context.Context) (*gosql.DB, error) { Password: password, User: username, }) - if err != nil { - return nil, err - } - db, err := gosql.Open("snowflake", dsn) - if err != nil { - return nil, err - } - return db, nil } // getSFCreds gets the snowflake credentials from the secrets manager diff --git a/pkg/cmd/roachtest/testselector/selector_test.go b/pkg/cmd/roachtest/testselector/selector_test.go new file mode 100644 index 000000000000..8b0179dc5ac2 --- /dev/null +++ b/pkg/cmd/roachtest/testselector/selector_test.go @@ -0,0 +1,159 @@ +// Copyright 2024 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 testselector + +import ( + "context" + gosql "database/sql" + "fmt" + "os" + "regexp" + "strings" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec" + "github.com/stretchr/testify/require" +) + +func TestCategoriseTests(t *testing.T) { + _ = os.Unsetenv(sfUsernameEnv) + _ = os.Unsetenv(sfPasswordEnv) + t.Run("expect getConnect to fail due to missing SFUSER env", func(t *testing.T) { + SqlConnectorFunc = nil + tds, err := CategoriseTests(context.Background(), nil) + require.Nil(t, tds) + require.NotNil(t, err) + require.Equal(t, "environment variable SFUSER is not set", err.Error()) + }) + _ = os.Setenv(sfUsernameEnv, "dummy_user") + _ = os.Setenv(sfPasswordEnv, "dummy_password") + t.Run("expect sql connector to fail", func(t *testing.T) { + SqlConnectorFunc = func(_, _ string) (*gosql.DB, error) { + return nil, fmt.Errorf("failed to connect to DB") + } + tds, err := CategoriseTests(context.Background(), nil) + require.Nil(t, tds) + require.NotNil(t, err) + require.Equal(t, "failed to connect to DB", err.Error()) + }) + var mock sqlmock.Sqlmock + var db *gosql.DB + SqlConnectorFunc = func(driverName, dataSourceName string) (*gosql.DB, error) { + dsn, err := getDSN() + require.Equal(t, "snowflake", driverName) + require.Equal(t, dsn, dataSourceName) + return db, err + } + var err error + t.Run("expect Prepare to fail", func(t *testing.T) { + db, mock, err = sqlmock.New() + require.Nil(t, err) + mock.ExpectPrepare(regexp.QuoteMeta(PreparedQuery)).WillReturnError(fmt.Errorf("failed to prepare")) + tds, err := CategoriseTests(context.Background(), nil) + require.Nil(t, tds) + require.NotNil(t, err) + require.Equal(t, "failed to prepare", err.Error()) + }) + t.Run("expect query to fail", func(t *testing.T) { + db, mock, err = sqlmock.New() + require.Nil(t, err) + mock.ExpectPrepare(regexp.QuoteMeta(PreparedQuery)) + mock.ExpectQuery(regexp.QuoteMeta(PreparedQuery)).WillReturnError(fmt.Errorf("failed to execute query")) + tds, err := CategoriseTests(context.Background(), &SelectTestsReq{ + ForPastDays: 1, + FirstRunOn: 2, + LastRunOn: 3, + Cloud: spec.AWS, + Suite: "unit test", + }) + require.Nil(t, tds) + require.NotNil(t, err) + require.Equal(t, "failed to execute query", err.Error()) + }) + t.Run("expect the sequence of response list is maintained for success", func(t *testing.T) { + db, mock, err = sqlmock.New() + mock.ExpectPrepare(regexp.QuoteMeta(PreparedQuery)) + rows := sqlmock.NewRows([]string{ + "name", "selected", "avg_duration", "last_failure_is_preempt", + }) + data := [][]string{ + {"t1", "no", "12345", "no"}, + {"t2", "no", "12345", "no"}, + {"t3", "no", "12345", "yes"}, + {"t4", "no", "12345", "no"}, + {"t5", "yes", "12345", "no"}, + {"t6", "yes", "12345", "no"}, + {"t7", "yes", "12345", "no"}, + {"t8", "yes", "12345", "no"}, + {"t9", "yes", "12345", "no"}, + } + for _, ds := range data { + rows.FromCSVString(strings.Join(ds, ",")) + } + mock.ExpectQuery(regexp.QuoteMeta(PreparedQuery)).WillReturnRows(rows) + tds, err := CategoriseTests(context.Background(), &SelectTestsReq{ + ForPastDays: 1, + FirstRunOn: 2, + LastRunOn: 3, + Cloud: spec.AWS, + Suite: "unit test", + }) + require.NotNil(t, tds) + require.Nil(t, err) + require.Equal(t, len(data), len(tds)) + // the sequence of response list must be maintained. + for i, d := range data { + td := tds[i] + require.Equal(t, d[0], td.Name) + require.Equal(t, d[1] != "no", td.Selected) + require.Equal(t, getDuration(d[2]), td.AvgDurationInMillis) + require.Equal(t, d[3] == "yes", td.LastFailureIsPreempt) + } + }) +} + +func TestNewDefaultSelectTestsReq(t *testing.T) { + req := NewDefaultSelectTestsReq(spec.Azure, "ut suite") + require.Equal(t, spec.Azure, req.Cloud) + require.Equal(t, "ut suite", req.Suite) + require.Equal(t, defaultForPastDays, req.ForPastDays) + require.Equal(t, defaultFirstRunOn, req.FirstRunOn) + require.Equal(t, defaultLastRunOn, req.LastRunOn) +} + +func Test_getSFCreds(t *testing.T) { + _ = os.Unsetenv(sfUsernameEnv) + _ = os.Unsetenv(sfPasswordEnv) + t.Run("expect username env failure", func(t *testing.T) { + u, p, e := getSFCreds() + require.Empty(t, u) + require.Empty(t, p) + require.NotNil(t, e) + require.Equal(t, fmt.Sprintf("environment variable %s is not set", sfUsernameEnv), e.Error()) + }) + t.Run("expect password env failure", func(t *testing.T) { + _ = os.Setenv(sfUsernameEnv, "dummy_user") + u, p, e := getSFCreds() + require.Empty(t, u) + require.Empty(t, p) + require.NotNil(t, e) + require.Equal(t, fmt.Sprintf("environment variable %s is not set", sfPasswordEnv), e.Error()) + }) + t.Run("expect no failure", func(t *testing.T) { + _ = os.Setenv(sfUsernameEnv, "dummy_user") + _ = os.Setenv(sfPasswordEnv, "dummy_password") + u, p, e := getSFCreds() + require.Equal(t, os.Getenv(sfUsernameEnv), u) + require.Equal(t, os.Getenv(sfPasswordEnv), p) + require.Nil(t, e) + }) +}