fix: out of range error + use local variables (#4)

Author: Eun

main.go

@@ -14,25 +14,22 @@ import (
 	"unsafe"
 )
 
-var bucketIndex *int
-var bucketCount *int
-var directoriesToExclude *string
-var packagesToExclude *string
-
-var directoriesToExcludeList []string
-var packagesToExcludeList []string
-
 // Buckets must be called to get the test bucket feature working.
 // It will modify the tests present in the testing.M struct.
 func Buckets(m *testing.M) {
+	var bucketIndex int
+	var bucketCount int
+
+	var directoriesToExclude []string
+	var packagesToExclude []string
+
 	if v := os.Getenv("BUCKET"); v != "" {
 		//nolint: gomnd // use 64 bits for parsing
 		n, err := strconv.ParseInt(v, 0, 64)
 		if err != nil {
 			panic(fmt.Sprintf("unable to parse BUCKET %s: %v", v, err))
 		}
-		i := int(n)
-		bucketIndex = &i
+		bucketIndex = int(n)
 	}
 
 	if v := os.Getenv("TOTAL_BUCKETS"); v != "" {
@@ -41,59 +38,33 @@ func Buckets(m *testing.M) {
 		if err != nil {
 			panic(fmt.Sprintf("unable to parse BUCKET_COUNT %s: %v", v, err))
 		}
-		i := int(n)
-		bucketCount = &i
+		bucketCount = int(n)
 	}
 
 	if v := os.Getenv("EXCLUDE_DIRECTORIES"); v != "" {
-		directoriesToExclude = &v
-	}
-
-	if v := os.Getenv("EXCLUDE_PACKAGES"); v != "" {
-		packagesToExclude = &v
-	}
-
-	if directoriesToExclude != nil {
-		directoriesToExcludeList = strings.FieldsFunc(*directoriesToExclude, func(r rune) bool {
+		directoriesToExclude = strings.FieldsFunc(v, func(r rune) bool {
 			return r == ',' || r == ';'
 		})
-		for i := range directoriesToExcludeList {
-			directoriesToExcludeList[i] = filepath.ToSlash(directoriesToExcludeList[i])
+		for i := range directoriesToExclude {
+			directoriesToExclude[i] = filepath.ToSlash(directoriesToExclude[i])
 		}
 	}
-	if packagesToExclude != nil {
-		packagesToExcludeList = strings.FieldsFunc(*packagesToExclude, func(r rune) bool {
+
+	if v := os.Getenv("EXCLUDE_PACKAGES"); v != "" {
+		packagesToExclude = strings.FieldsFunc(v, func(r rune) bool {
 			return r == ',' || r == ';'
 		})
 	}
 
-	if !isBucketFeatureEnabled() && !isDirectoriesToExcludeEnabled() && !isPackagesToExcludeEnabled() {
+	if (bucketCount == 0 || bucketIndex >= bucketCount) && (len(directoriesToExclude) == 0 && len(packagesToExclude) == 0) {
 		return
 	}
 
 	v := reflect.ValueOf(m).Elem()
 	testsField := v.FieldByName("tests")
 	//nolint: gosec // allow the usage of unsafe so we can get the test slice.
 	ptr := unsafe.Pointer(testsField.UnsafeAddr())
-	filterTests((*[]testing.InternalTest)(ptr))
-}
-func isBucketFeatureEnabled() bool {
-	if bucketCount == nil || bucketIndex == nil {
-		return false
-	}
-
-	if *bucketCount <= 0 || *bucketIndex < 0 || *bucketIndex >= *bucketCount {
-		return false
-	}
-	return true
-}
-
-func isDirectoriesToExcludeEnabled() bool {
-	return len(directoriesToExcludeList) > 0
-}
-
-func isPackagesToExcludeEnabled() bool {
-	return len(packagesToExcludeList) > 0
+	filterTests((*[]testing.InternalTest)(ptr), bucketIndex, bucketCount, directoriesToExclude, packagesToExclude)
 }
 
 func getSourceFile(f func(*testing.T)) string {
@@ -153,39 +124,40 @@ dirLoop:
 	return false
 }
 
-func filterTests(tests *[]testing.InternalTest) {
-	if isDirectoriesToExcludeEnabled() {
+func filterTests(tests *[]testing.InternalTest, bucketIndex, bucketCount int, directoriesToExclude, packagesToExclude []string) {
+	if len(directoriesToExclude) > 0 {
 		for i := len(*tests) - 1; i >= 0; i-- {
 			file := getSourceFile((*tests)[i].F)
 			if file == "" {
 				fmt.Printf("unable to find source of %s\n", (*tests)[i].Name)
 				continue
 			}
 
-			if isFileInDir(filepath.ToSlash(file), directoriesToExcludeList...) {
+			if isFileInDir(filepath.ToSlash(file), directoriesToExclude...) {
 				*tests = append((*tests)[:i], (*tests)[i+1:]...)
 			}
 		}
 	}
-	if isPackagesToExcludeEnabled() {
+	if len(packagesToExclude) > 0 {
 		for i := len(*tests) - 1; i >= 0; i-- {
 			pkg := getPackageName((*tests)[i].F)
 			if pkg == "" {
 				fmt.Printf("unable to find package of %s\n", (*tests)[i].Name)
 				continue
 			}
-			if isFileInDir(pkg, packagesToExcludeList...) {
+			if isFileInDir(pkg, packagesToExclude...) {
 				*tests = append((*tests)[:i], (*tests)[i+1:]...)
 			}
 		}
 	}
-	if isBucketFeatureEnabled() {
-		perBucket := int(math.Ceil(float64(len(*tests)) / float64(*bucketCount)))
 
-		from := *bucketIndex * perBucket
+	if bucketCount > 0 && bucketIndex >= 0 && bucketIndex < bucketCount {
+		perBucket := int(math.Ceil(float64(len(*tests)) / float64(bucketCount)))
+
+		from := bucketIndex * perBucket
 		to := from + perBucket
 
-		if to > len(*tests) {
+		if to > len(*tests)-1 {
 			to = len(*tests)
 		}
 

main_test.go

@@ -106,7 +106,7 @@ func subTest(t *testing.T, env []string, coverProfile, packageName string, testS
 	cmd.Env = append(os.Environ(), env...)
 
 	if err := cmd.Run(); err != nil {
-		t.Fatal(err)
+		t.Fatal(err, buf.String())
 	}
 
 	type test struct {
@@ -163,3 +163,77 @@ func subTest(t *testing.T, env []string, coverProfile, packageName string, testS
 		}
 	}
 }
+
+func Test_filterTestsBuckets(t *testing.T) {
+	testCases := []struct {
+		name        string
+		testCount   int
+		bucketCount int
+		buckets     []int
+	}{
+		{
+			"1/1",
+			1,
+			1,
+			[]int{1},
+		},
+		{
+			"1/2",
+			1,
+			2,
+			[]int{1, 0},
+		},
+		{
+			"3/2",
+			3,
+			2,
+			[]int{2, 1},
+		},
+		{
+			"4/2",
+			4,
+			2,
+			[]int{2, 2},
+		},
+		{
+			"5/2",
+			5,
+			2,
+			[]int{3, 2},
+		},
+		{
+			"11/3",
+			11,
+			3,
+			[]int{4, 4, 3},
+		},
+		{
+			"12/3",
+			12,
+			3,
+			[]int{4, 4, 4},
+		},
+		{
+			"13/3",
+			13,
+			3,
+			[]int{5, 5, 3},
+		},
+	}
+	for _, tt := range testCases {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			for i := 0; i < tt.bucketCount; i++ {
+				tests := make([]testing.InternalTest, tt.testCount)
+				for j := 0; j < tt.testCount; j++ {
+					tests[j].Name = fmt.Sprintf("Test%d", j)
+				}
+				filterTests(&tests, i, tt.bucketCount, nil, nil)
+
+				if tt.buckets[i] != len(tests) {
+					t.Fatalf("expected %d tests in bucket %d but got %d", tt.buckets[i], i, len(tests))
+				}
+			}
+		})
+	}
+}