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

@@ -93,6 +93,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

@@ -163,6 +163,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",
@@ -1219,6 +1220,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",
@@ -124,6 +125,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,102 @@ 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", func(t *testing.T) {
+		db, mock, err = sqlmock.New()
+		require.Nil(t, err)
+		specs, _ := getMockData()
+		mock.ExpectPrepare(regexp.QuoteMeta(testselector.PreparedQuery)).WillReturnError(fmt.Errorf("failed to prepare"))
+		updateSpecForSelectiveTests(ctx, specs)
+		for _, spec := range specs {
+			if !strings.Contains(spec.Name, "skipped") {
+				require.Empty(t, spec.Skip)
+			} else {
+				require.Equal(t, "dummy skip", spec.Skip)
+			}
+		}
+	})
+	t.Run("expect no failure", func(t *testing.T) {
+		db, mock, err = sqlmock.New()
+		require.Nil(t, err)
+		specs, rows := getMockData()
+		mock.ExpectPrepare(regexp.QuoteMeta(testselector.PreparedQuery))
+		mock.ExpectQuery(regexp.QuoteMeta(testselector.PreparedQuery)).WillReturnRows(rows)
+		oldSuite := roachtestflags.Suite
+		roachtestflags.Suite = registry.Nightly
+		roachtestflags.SuccessfulTestsSelectPct = 0.30
+		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, "dummy skip", s.Skip)
+				require.Equal(t, "dummy skip test", s.SkipDetails)
+			} else {
+				require.Empty(t, s.Skip)
+			}
+		}
+		roachtestflags.Suite = oldSuite
+	})
+}
+
+func getMockData() ([]registry.TestSpec, *sqlmock.Rows) {
+	specs := []registry.TestSpec{
+		{Name: "t1_selected"},
+		{Name: "t2_selected"},
+		{Name: "t_skipped_selected", Skip: "dummy skip", SkipDetails: "dummy skip test"},
+		{Name: "t1_success"},
+		{Name: "t_randomized", Randomized: true},
+		{Name: "t2_success"},
+		{Name: "t3_success"},
+		{Name: "t5_success_skip_selector"},
+		{Name: "t_opt_out", TestSelectionOptOutSuites: registry.Suites(registry.Nightly)},
+		{Name: "t6_success_skip_selector", TestSelectionOptOutSuites: registry.Suites(registry.Weekly)},
+		{Name: "t_skipped_not_selected", Skip: "dummy skip", SkipDetails: "dummy skip test"},
+		{Name: "t8_success_skip_selector"},
+		{Name: "t10_success_skip_selector"},
+		{Name: "t10_success_skip_selector"},
+	}
+	for _, s := range specs {
+		s.Suites = registry.Suites(registry.Nightly, registry.Weekly)
+	}
+	data := [][]string{
+		{"t1_selected", "yes", "1", "no"},
+		{"t2_selected", "yes", "2", "yes"},
+		{"t_skipped_selected", "yes", "3", "no"},
+		{"t1_success", "no", "4", "no"},
+		{"t_ramdomized", "no", "5", "no"},
+		{"t2_success", "no", "6", "no"},
+		{"t3_success", "no", "7", "no"},
+		{"t_opt_out", "no", "8", "no"},
+		{"t4_missing", "no", "9", "no"},
+		{"t5_success_skip_selector", "no", "10", "no"},
+		{"t6_success_skip_selector", "no", "11", "no"},
+		{"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"},
+	}
+	rows := sqlmock.NewRows([]string{
+		"name", "selected", "avg_duration", "last_failure_is_preempt",
+	})
+
+	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) (map[string]*Test
 		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) (map[string]*Test
 	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