roachtest: UT for test selection

This PR has the unit tests for the test selection
changes. The tests the selector and the main code.
This does not validate the snowflake query, but,
simulates the response for other tests.

Fixes: #129786
Epic: None

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

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",

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",

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
+}

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",
+    ],
+)

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,28 +151,33 @@ 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,
 		Warehouse: warehouse,
 		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