Repository: pingcap/tidb (go)

Issue:
provide `intest.AssertNotNil` and `intest.AssertNoError`
## Enhancement

`intest.Assert` receives an `any` parameter as the condition. Sometimes, it is misleading. For example, if the user wants to check an error is nil, he should write the code like this:

```go
err := doSomething()
intest.Assert(err != nil)
```

The above code is the right practice, but because of the input parameter is `any`, he can also write like this:

```go
intest.Assert(err) // it's a mistake, the code will panic when err is nil!
```

It's better to make Assert only receives a bool input and provides `AssertNotNil` and `AssertNoError` to assert a pointer is not nil or an error is nil.



Pull request description:
util/intest: provide `intest.AssertNotNil` and `intest.AssertNoError`
<!--

Thank you for contributing to TiDB!

PR Title Format:
1. pkg [, pkg2, pkg3]: what's changed
2. *: what's changed

-->

### What problem does this PR solve?
<!--

Please create an issue first to describe the problem.

There MUST be one line starting with "Issue Number:  " and 
linking the relevant issues via the "close" or "ref".

For more info, check https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/contribute-code.html#referring-to-an-issue.

-->

Issue Number: close #48578

### What is changed and how it works?

Make `intest.Assert` to have a brief semantics:

- `intest.Assert` only receives `bool` input to check the condition should be true
- provide `intest.AssertNotNil` to check a pointer should not be nil
- provide `intest.AssertNoError` to check a error should be nil 

### Check List

Tests <!-- At least one of them must be included. -->

- [x] Unit test
- [ ] Integration test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test
  > - [ ] I checked and no code files have been changed.
  > <!-- Or your custom  "No need to test" reasons -->  

Side effects

- [ ] Performance regression: Consumes more CPU
- [ ] Performance regression: Consumes more Memory
- [ ] Breaking backward compatibility

Documentation

- [ ] Affects user behaviors
- [ ] Contains syntax changes
- [ ] Contains variable changes
- [ ] Contains experimental features
- [ ] Changes MySQL compatibility

### Release note

<!-- compatibility change, improvement, bugfix, and new feature need a release note -->

Please refer to [Release Notes Language Style Guide](https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/release-notes-style-guide.html) to write a quality release note.

```release-note
None
```


Patch:
diff --git a/pkg/sessionctx/stmtctx/stmtctx.go b/pkg/sessionctx/stmtctx/stmtctx.go
index 2216d05ef3..d7fa9eed2f 100644
--- a/pkg/sessionctx/stmtctx/stmtctx.go
+++ b/pkg/sessionctx/stmtctx/stmtctx.go
@@ -431,7 +431,7 @@ func NewStmtCtx() *StatementContext {
 
 // NewStmtCtxWithTimeZone creates a new StatementContext with the given timezone
 func NewStmtCtxWithTimeZone(tz *time.Location) *StatementContext {
-	intest.Assert(tz)
+	intest.AssertNotNil(tz)
 	sc := &StatementContext{}
 	sc.typeCtx = types.NewContext(types.DefaultStmtFlags, tz, sc.AppendWarning)
 	return sc
@@ -451,7 +451,7 @@ func (sc *StatementContext) TimeZone() *time.Location {
 
 // SetTimeZone sets the timezone
 func (sc *StatementContext) SetTimeZone(tz *time.Location) {
-	intest.Assert(tz)
+	intest.AssertNotNil(tz)
 	sc.typeCtx = sc.typeCtx.WithLocation(tz)
 }
 
diff --git a/pkg/types/context.go b/pkg/types/context.go
index b969e4801f..af678ebc67 100644
--- a/pkg/types/context.go
+++ b/pkg/types/context.go
@@ -211,7 +211,7 @@ func (c *Context) WithFlags(f Flags) Context {
 
 // WithLocation returns a new context with the given location
 func (c *Context) WithLocation(loc *time.Location) Context {
-	intest.Assert(loc)
+	intest.AssertNotNil(loc)
 	ctx := *c
 	ctx.loc = loc
 	return ctx
@@ -219,7 +219,7 @@ func (c *Context) WithLocation(loc *time.Location) Context {
 
 // Location returns the location of the context
 func (c *Context) Location() *time.Location {
-	intest.Assert(c.loc)
+	intest.AssertNotNil(c.loc)
 	if c.loc == nil {
 		// c.loc should always not be nil, just make the code safe here.
 		return time.UTC


Test patch:
diff --git a/pkg/util/intest/BUILD.bazel b/pkg/util/intest/BUILD.bazel
index bab02008d0..d34d4b1cbd 100644
--- a/pkg/util/intest/BUILD.bazel
+++ b/pkg/util/intest/BUILD.bazel
@@ -18,6 +18,7 @@ go_test(
     flaky = True,
     deps = [
         ":intest",
+        "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )
diff --git a/pkg/util/intest/assert.go b/pkg/util/intest/assert.go
index c40d6cc71a..1d96c38ff1 100644
--- a/pkg/util/intest/assert.go
+++ b/pkg/util/intest/assert.go
@@ -21,63 +21,66 @@ import (
 	"reflect"
 )
 
-// Assert asserts a condition. It only works in test (intest.InTest == true).
-// You can assert a condition like this to assert a variable `foo` is not nil: `assert.Assert(foo != nil)`.
-// Or you can pass foo as a parameter directly for simple: `assert.Assert(foo)`
-// You can also assert a function that returns a bool: `intest.Assert(func() bool { return foo != nil })`
-// If you pass a function without a signature `func() bool`, the function will always panic.
-func Assert(cond any, msgAndArgs ...any) {
+// Assert asserts a condition is true
+func Assert(cond bool, msgAndArgs ...any) {
+	if InTest && !cond {
+		doPanic("", msgAndArgs...)
+	}
+}
+
+// AssertNoError asserts an error is nil
+func AssertNoError(err error, msgAndArgs ...any) {
+	if InTest && err != nil {
+		doPanic(fmt.Sprintf("error is not nil: %+v", err), msgAndArgs...)
+	}
+}
+
+// AssertNotNil asserts an object is not nil
+func AssertNotNil(obj any, msgAndArgs ...any) {
 	if InTest {
-		assert(cond, msgAndArgs...)
+		Assert(obj != nil, msgAndArgs...)
+		value := reflect.ValueOf(obj)
+		switch value.Kind() {
+		case reflect.Func, reflect.Chan, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice, reflect.UnsafePointer:
+			Assert(!value.IsNil(), msgAndArgs...)
+		}
 	}
 }
 
 // AssertFunc asserts a function condition
 func AssertFunc(fn func() bool, msgAndArgs ...any) {
 	if InTest {
-		assert(fn(), msgAndArgs...)
+		Assert(fn != nil, msgAndArgs...)
+		Assert(fn(), msgAndArgs...)
 	}
 }
 
-func assert(cond any, msgAndArgs ...any) {
-	if !checkAssertObject(cond) {
-		doPanic(msgAndArgs...)
-	}
+func doPanic(extraMsg string, userMsgAndArgs ...any) {
+	panic(assertionFailedMsg(extraMsg, userMsgAndArgs...))
 }
 
-func doPanic(msgAndArgs ...any) {
-	panic(assertionFailedMsg(msgAndArgs...))
-}
+func assertionFailedMsg(extraMsg string, userMsgAndArgs ...any) string {
+	msg := "assert failed"
+	if len(userMsgAndArgs) == 0 {
+		if extraMsg != "" {
+			msg = fmt.Sprintf("%s, %s", msg, extraMsg)
+		}
+		return msg
+	}
 
-func assertionFailedMsg(msgAndArgs ...any) string {
-	if len(msgAndArgs) == 0 {
-		return "assert failed"
+	if len(userMsgAndArgs) == 0 {
+		return fmt.Sprintf("assert failed, %s", extraMsg)
 	}
 
-	msg, ok := msgAndArgs[0].(string)
+	userMsg, ok := userMsgAndArgs[0].(string)
 	if !ok {
-		msg = fmt.Sprintf("%+v", msgAndArgs[0])
+		userMsg = fmt.Sprintf("%+v", userMsgAndArgs[0])
 	}
 
-	msg = fmt.Sprintf("assert failed: %s", msg)
-	return fmt.Sprintf(msg, msgAndArgs[1:]...)
-}
-
-func checkAssertObject(obj any) bool {
-	if obj == nil {
-		return false
+	msg = fmt.Sprintf("%s, %s", msg, userMsg)
+	if extraMsg != "" {
+		msg = fmt.Sprintf("%s, %s", msg, extraMsg)
 	}
 
-	value := reflect.ValueOf(obj)
-	switch value.Kind() {
-	case reflect.Bool:
-		return obj.(bool)
-	case reflect.Func:
-		panic("you should use `intest.Assert(fn != nil)` to assert a function is not nil, " +
-			"or use `intest.AssertFunc(fn)` to assert a function's return value is true")
-	case reflect.Chan, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice, reflect.UnsafePointer:
-		return !value.IsNil()
-	default:
-		return true
-	}
+	return fmt.Sprintf(msg, userMsgAndArgs[1:]...)
 }
diff --git a/pkg/util/intest/assert_test.go b/pkg/util/intest/assert_test.go
index c85fc8d5b1..6a14cf7939 100644
--- a/pkg/util/intest/assert_test.go
+++ b/pkg/util/intest/assert_test.go
@@ -15,8 +15,10 @@
 package intest_test
 
 import (
+	"strings"
 	"testing"
 
+	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/stretchr/testify/require"
 )
@@ -26,37 +28,55 @@ type foo struct{}
 func TestAssert(t *testing.T) {
 	require.True(t, intest.InTest)
 	checkAssert(t, true, true)
-	checkAssert(t, "", true)
-	checkAssert(t, "abc", true)
-	checkAssert(t, 0, true)
-	checkAssert(t, 123, true)
-	checkAssert(t, foo{}, true)
-	checkAssert(t, &foo{}, true)
 	checkAssert(t, false, false)
-	checkAssert(t, false, false, "assert failed: msg1", "msg1")
-	checkAssert(t, false, false, "assert failed: msg2 a b 1", "msg2 %s %s %d", "a", "b", 1)
-	checkAssert(t, false, false, "assert failed: 123", 123)
-	checkAssert(t, nil, false)
+	checkAssert(t, false, false, "assert failed, msg1", "msg1")
+	checkAssert(t, false, false, "assert failed, msg2 a b 1", "msg2 %s %s %d", "a", "b", 1)
+	checkAssert(t, false, false, "assert failed, 123", 123)
+	checkAssertNotNil(t, "", true)
+	checkAssertNotNil(t, "abc", true)
+	checkAssertNotNil(t, 0, true)
+	checkAssertNotNil(t, 123, true)
+	checkAssertNotNil(t, foo{}, true)
+	checkAssertNotNil(t, &foo{}, true)
+	checkAssertNotNil(t, nil, false)
+	checkAssertNotNil(t, true, true)
+	checkAssertNotNil(t, false, true)
 	var f *foo
-	checkAssert(t, f, false)
-	checkAssert(t, func() bool { return true }, false, "you should use `intest.Assert(fn != nil)` to assert a function is not nil, or use `intest.AssertFunc(fn)` to assert a function's return value is true")
-	checkAssert(t, func(_ string) bool { return true }, false, "you should use `intest.Assert(fn != nil)` to assert a function is not nil, or use `intest.AssertFunc(fn)` to assert a function's return value is true")
+	checkAssertNotNil(t, f, false)
+	checkAssertNotNil(t, func() bool { return true }, true)
+	checkAssertNotNil(t, func() bool { return false }, true)
+	checkAssertNotNil(t, func(_ string) bool { return true }, true)
+	checkAssertNotNil(t, nil, false, "assert failed, msg1", "msg1")
+	checkAssertNotNil(t, nil, false, "assert failed, msg2 a b 1", "msg2 %s %s %d", "a", "b", 1)
 	checkFuncAssert(t, func() bool { panic("inner panic1") }, false, "inner panic1")
 	checkFuncAssert(t, func() bool { return true }, true)
 	checkFuncAssert(t, func() bool { return false }, false)
-	checkFuncAssert(t, func() bool { return false }, false, "assert failed: msg3", "msg3")
-	checkFuncAssert(t, func() bool { return false }, false, "assert failed: msg4 c d 2", "msg4 %s %s %d", "c", "d", 2)
+	checkFuncAssert(t, func() bool { return false }, false, "assert failed, msg3", "msg3")
+	checkFuncAssert(t, func() bool { return false }, false, "assert failed, msg4 c d 2", "msg4 %s %s %d", "c", "d", 2)
 	checkFuncAssert(t, func() bool { panic("inner panic2") }, false, "inner panic2")
+	checkFuncAssert(t, nil, false)
+	checkAssertNoError(t, nil, true)
+	checkAssertNoError(t, errors.New("mock err1"), false, "assert failed, error is not nil: mock err1")
+	var err error
+	checkAssertNoError(t, err, true)
 }
 
 func checkFuncAssert(t *testing.T, fn func() bool, pass bool, msgAndArgs ...any) {
 	doCheckAssert(t, intest.AssertFunc, fn, pass, msgAndArgs...)
 }
 
-func checkAssert(t *testing.T, cond any, pass bool, msgAndArgs ...any) {
+func checkAssert(t *testing.T, cond bool, pass bool, msgAndArgs ...any) {
 	doCheckAssert(t, intest.Assert, cond, pass, msgAndArgs...)
 }
 
+func checkAssertNotNil(t *testing.T, obj any, pass bool, msgAndArgs ...any) {
+	doCheckAssert(t, intest.AssertNotNil, obj, pass, msgAndArgs...)
+}
+
+func checkAssertNoError(t *testing.T, err error, pass bool, msgAndArgs ...any) {
+	doCheckAssert(t, intest.AssertNoError, err, pass, msgAndArgs...)
+}
+
 func doCheckAssert(t *testing.T, fn any, cond any, pass bool, msgAndArgs ...any) {
 	expectMsg := "assert failed"
 	if len(msgAndArgs) > 0 {
@@ -64,32 +84,39 @@ func doCheckAssert(t *testing.T, fn any, cond any, pass bool, msgAndArgs ...any)
 		msgAndArgs = msgAndArgs[1:]
 	}
 
+	fail := ""
+	onlyCheckPrefix := false
 	if !pass {
 		defer func() {
+			if fail != "" {
+				require.FailNow(t, fail)
+			}
 			r := recover()
 			require.NotNil(t, r)
-			require.Equal(t, expectMsg, r)
+			if onlyCheckPrefix {
+				require.True(t, strings.HasPrefix(r.(string), expectMsg), "%q\nshould have prefix \n%q", r, expectMsg)
+				require.Contains(t, r.(string), expectMsg)
+			} else {
+				require.Equal(t, expectMsg, r)
+			}
 		}()
 	}
 
-	testFn, ok := fn.(func(any, ...any))
-	if !ok {
-		if fnAssert, ok := fn.(func(func() bool, ...any)); ok {
-			testFn = func(any, ...any) {
-				fnAssert(cond.(func() bool), msgAndArgs...)
-			}
+	switch assertFn := fn.(type) {
+	case func(bool, ...any):
+		assertFn(cond.(bool), msgAndArgs...)
+	case func(any, ...any):
+		assertFn(cond, msgAndArgs...)
+	case func(func() bool, ...any):
+		assertFn(cond.(func() bool), msgAndArgs...)
+	case func(error, ...any):
+		if cond == nil {
+			assertFn(nil, msgAndArgs...)
 		} else {
-			require.FailNow(t, "invalid assert function")
+			onlyCheckPrefix = true
+			assertFn(cond.(error), msgAndArgs...)
 		}
+	default:
+		fail = "invalid input assert function"
 	}
-
-	if len(msgAndArgs) == 0 {
-		testFn(cond)
-	}
-
-	if len(msgAndArgs) == 1 {
-		testFn(cond, msgAndArgs[0])
-	}
-
-	testFn(cond, msgAndArgs...)
 }
diff --git a/pkg/util/intest/common.go b/pkg/util/intest/common.go
index 56b928bc9e..a2ebccd4b3 100644
--- a/pkg/util/intest/common.go
+++ b/pkg/util/intest/common.go
@@ -21,7 +21,15 @@ const InTest = false
 
 // Assert is a stub function in release build.
 // See the same function in `util/intest/assert.go` for the real implement in test.
-func Assert(_ any, _ ...any) {}
+func Assert(_ bool, _ ...any) {}
+
+// AssertNotNil is a stub function in release build.
+// See the same function in `util/intest/assert.go` for the real implement in test.
+func AssertNotNil(_ any, _ ...any) {}
+
+// AssertNoError is a stub function in release build.
+// See the same function in `util/intest/assert.go` for the real implement in test.
+func AssertNoError(_ error, _ ...any) {}
 
 // AssertFunc is a stub function in release build.
 // See the same function `util/intest/assert.go` for the real implement in test.