From 23d79d48c2eedb0fa773dda2609e944f6e111505 Mon Sep 17 00:00:00 2001
From: tangenta <tangenta@126.com>
Date: Wed, 23 Mar 2022 11:38:33 +0800
Subject: [PATCH] tools/check/ut: support pattern matching on test names
 (#33020)

ref pingcap/tidb#30822
---
 tools/check/ut.go | 52 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 8 deletions(-)

diff --git a/tools/check/ut.go b/tools/check/ut.go
index 5e47cf13c1866..4f8bab7d8cce3 100644
--- a/tools/check/ut.go
+++ b/tools/check/ut.go
@@ -24,6 +24,7 @@ import (
 	"os"
 	"os/exec"
 	"path"
+	"regexp"
 	"runtime"
 	"sort"
 	"strings"
@@ -48,6 +49,9 @@ ut list
 // list test cases of a single package
 ut list $package
 
+// list test cases that match a pattern
+ut list $package 'r:$regex'
+
 // run all tests
 ut run
 
@@ -57,6 +61,9 @@ ut run $package
 // run test cases of a single package
 ut run $package $test
 
+// run test cases that match a pattern
+ut run $package 'r:$regex'
+
 // build all test package
 ut build
 
@@ -101,7 +108,7 @@ func cmdList(args ...string) bool {
 	}
 
 	// list test case of a single package
-	if len(args) == 1 {
+	if len(args) == 1 || len(args) == 2 {
 		pkg := args[0]
 		pkgs = filter(pkgs, func(s string) bool { return s == pkg })
 		if len(pkgs) != 1 {
@@ -129,6 +136,15 @@ func cmdList(args ...string) bool {
 			fmt.Println("list test cases for package error", err)
 			return false
 		}
+
+		if len(args) == 2 {
+			res, err = filterTestCases(res, args[1])
+			if err != nil {
+				fmt.Println("filter test cases error", err)
+				return false
+			}
+		}
+
 		for _, x := range res {
 			fmt.Println(x.test)
 		}
@@ -249,14 +265,11 @@ func cmdRun(args ...string) bool {
 			fmt.Println("list test cases error", err)
 			return false
 		}
-		// filter the test case to run
-		tmp := tasks[:0]
-		for _, task := range tasks {
-			if strings.Contains(task.test, args[1]) {
-				tmp = append(tmp, task)
-			}
+		tasks, err = filterTestCases(tasks, args[1])
+		if err != nil {
+			fmt.Println("filter test cases error", err)
+			return false
 		}
-		tasks = tmp
 	}
 
 	if except != "" {
@@ -608,6 +621,29 @@ func listTestCases(pkg string, tasks []task) ([]task, error) {
 	return tasks, nil
 }
 
+func filterTestCases(tasks []task, arg1 string) ([]task, error) {
+	if strings.HasPrefix(arg1, "r:") {
+		r, err := regexp.Compile(arg1[2:])
+		if err != nil {
+			return nil, err
+		}
+		tmp := tasks[:0]
+		for _, task := range tasks {
+			if r.MatchString(task.test) {
+				tmp = append(tmp, task)
+			}
+		}
+		return tmp, nil
+	}
+	tmp := tasks[:0]
+	for _, task := range tasks {
+		if strings.Contains(task.test, arg1) {
+			tmp = append(tmp, task)
+		}
+	}
+	return tmp, nil
+}
+
 func listPackages() ([]string, error) {
 	cmd := exec.Command("go", "list", "./...")
 	ss, err := cmdToLines(cmd)