util: add InTest to check whether to be in test (#40927)

close #40929

Makefile

@@ -411,7 +411,7 @@ bazel_test: failpoint-enable bazel_ci_prepare
 
 bazel_coverage_test: failpoint-enable bazel_ci_prepare
 	bazel $(BAZEL_GLOBAL_CONFIG) coverage $(BAZEL_CMD_CONFIG) \
-		--build_event_json_file=bazel_1.json --@io_bazel_rules_go//go/config:cover_format=go_cover --define gotags=deadlock \
+		--build_event_json_file=bazel_1.json --@io_bazel_rules_go//go/config:cover_format=go_cover --define gotags=deadlock,intest \
 		-- //... -//cmd/... -//tests/graceshutdown/... \
 		-//tests/globalkilltest/... -//tests/readonlytest/... -//br/pkg/task:task_test -//tests/realtikvtest/...
 
@@ -444,27 +444,27 @@ bazel_golangcilinter:
 	-- run  $$($(PACKAGE_DIRECTORIES)) --config ./.golangci.yaml
 
 bazel_brietest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/brietest/...
 
 bazel_pessimistictest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/pessimistictest/...
 
 bazel_sessiontest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/sessiontest/...
 
 bazel_statisticstest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/statisticstest/...
 
 bazel_txntest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/txntest/...
 
 bazel_addindextest: failpoint-enable bazel_ci_prepare
-	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock \
+	bazel $(BAZEL_GLOBAL_CONFIG) test $(BAZEL_CMD_CONFIG) --test_arg=-with-real-tikv --define gotags=deadlock,intest \
 		-- //tests/realtikvtest/addindextest/...
 
 bazel_lint: bazel_prepare

ddl/BUILD.bazel

@@ -107,6 +107,7 @@ go_library(
         "//util/gcutil",
         "//util/generic",
         "//util/hack",
+        "//util/intest",
         "//util/logutil",
         "//util/mathutil",
         "//util/mock",

ddl/sanity_check.go

@@ -16,7 +16,6 @@ package ddl
 
 import (
 	"context"
-	"flag"
 	"fmt"
 	"strings"
 
@@ -26,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/pingcap/tidb/util/sqlexec"
@@ -182,7 +182,7 @@ func (ctx *delRangeCntCtx) deduplicateIdxCnt(indexIDs []int64) int {
 // It's only check during the test environment, so it would panic directly.
 // These checks may be controlled by configuration in the future.
 func (d *ddl) checkHistoryJobInTest(ctx sessionctx.Context, historyJob *model.Job) {
-	if !(flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil) {
+	if !intest.InTest {
 		return
 	}
 

kv/BUILD.bazel

@@ -30,6 +30,7 @@ go_library(
         "//types",
         "//util/codec",
         "//util/dbterror",
+        "//util/intest",
         "//util/logutil",
         "//util/memory",
         "//util/set",

kv/txn.go

@@ -17,7 +17,6 @@ package kv
 import (
 	"context"
 	"errors"
-	"flag"
 	"fmt"
 	"math"
 	"math/rand"
@@ -26,6 +25,7 @@ import (
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/parser/terror"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/tikv/client-go/v2/oracle"
 	"go.uber.org/zap"
@@ -206,7 +206,7 @@ func setRequestSourceForInnerTxn(ctx context.Context, txn Transaction) {
 	}
 	// panic in test mode in case there are requests without source in the future.
 	// log warnings in production mode.
-	if flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil {
+	if intest.InTest {
 		panic("unexpected no source type context, if you see this error, " +
 			"the `RequestSourceTypeKey` is missing in your context")
 	} else {

resourcemanager/util/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
     importpath = "github.com/pingcap/tidb/resourcemanager/util",
     visibility = ["//visibility:public"],
     deps = [
+        "//util/intest",
         "@com_github_pingcap_errors//:errors",
         "@org_uber_go_atomic//:atomic",
     ],
@@ -19,5 +20,8 @@ go_test(
     name = "util_test",
     srcs = ["shard_pool_map_test.go"],
     embed = [":util"],
-    deps = ["@com_github_stretchr_testify//require"],
+    deps = [
+        "//util/intest",
+        "@com_github_stretchr_testify//require",
+    ],
 )

resourcemanager/util/shard_pool_map.go

@@ -18,6 +18,7 @@ import (
 	"sync"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/util/intest"
 )
 
 const shard = 8
@@ -69,7 +70,7 @@ func newPoolMap() poolMap {
 func (p *poolMap) Add(key string, pool *PoolContainer) error {
 	p.mu.Lock()
 	defer p.mu.Unlock()
-	if _, contain := p.poolMap[key]; contain {
+	if _, contain := p.poolMap[key]; contain && !intest.InTest {
 		return errors.New("pool is already exist")
 	}
 	p.poolMap[key] = pool

resourcemanager/util/shard_pool_map_test.go

@@ -19,6 +19,7 @@ import (
 	"sync/atomic"
 	"testing"
 
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/stretchr/testify/require"
 )
 
@@ -29,7 +30,9 @@ func TestShardPoolMap(t *testing.T) {
 		id := strconv.FormatInt(int64(i), 10)
 		require.NoError(t, pm.Add(id, &PoolContainer{Pool: NewMockGPool(id), Component: DDL}))
 	}
-	require.Error(t, pm.Add("1", &PoolContainer{Pool: NewMockGPool("1"), Component: DDL}))
+	if !intest.InTest {
+		require.Error(t, pm.Add("1", &PoolContainer{Pool: NewMockGPool("1"), Component: DDL}))
+	}
 	var cnt atomic.Int32
 	pm.Iter(func(pool *PoolContainer) {
 		cnt.Add(1)

server/BUILD.bazel

@@ -78,6 +78,7 @@ go_library(
         "//util/fastrand",
         "//util/gcutil",
         "//util/hack",
+        "//util/intest",
         "//util/logutil",
         "//util/memory",
         "//util/pdapi",

server/mock_conn.go

@@ -18,7 +18,6 @@ import (
 	"bufio"
 	"bytes"
 	"context"
-	"flag"
 	"math/rand"
 	"testing"
 
@@ -29,6 +28,7 @@ import (
 	tmysql "github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/util/arena"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/stretchr/testify/require"
 )
 
@@ -80,7 +80,7 @@ func (mc *mockConn) ID() uint64 {
 func CreateMockServer(t *testing.T, store kv.Storage) *Server {
 	if !RunInGoTest {
 		// If CreateMockServer is called in another package, RunInGoTest is not initialized.
-		RunInGoTest = flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil
+		RunInGoTest = intest.InTest
 	}
 	tidbdrv := NewTiDBDriver(store)
 	cfg := config.NewConfig()

session/BUILD.bazel

@@ -72,6 +72,7 @@ go_library(
         "//util/collate",
         "//util/dbterror",
         "//util/execdetails",
+        "//util/intest",
         "//util/kvcache",
         "//util/logutil",
         "//util/logutil/consistency",

session/bootstrap.go

@@ -21,7 +21,6 @@ package session
 import (
 	"context"
 	"encoding/hex"
-	"flag"
 	"fmt"
 	"io/ioutil"
 	osuser "os/user"
@@ -51,6 +50,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/dbterror"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/pingcap/tidb/util/logutil"
 	utilparser "github.com/pingcap/tidb/util/parser"
 	"github.com/pingcap/tidb/util/sqlexec"
@@ -2407,12 +2407,6 @@ func doBootstrapSQLFile(s Session) {
 	}
 }
 
-// inTestSuite checks if we are bootstrapping in the context of tests.
-// There are some historical differences in behavior between tests and non-tests.
-func inTestSuite() bool {
-	return flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil
-}
-
 // doDMLWorks executes DML statements in bootstrap stage.
 // All the statements run in a single transaction.
 func doDMLWorks(s Session) {
@@ -2457,11 +2451,11 @@ func doDMLWorks(s Session) {
 				vVal = variable.On
 			}
 		case variable.TiDBMemOOMAction:
-			if inTestSuite() {
+			if intest.InTest {
 				vVal = variable.OOMActionLog
 			}
 		case variable.TiDBEnableAutoAnalyze:
-			if inTestSuite() {
+			if intest.InTest {
 				vVal = variable.Off
 			}
 		// For the following sysvars, we change the default

session/session.go

@@ -25,7 +25,6 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	stderrs "errors"
-	"flag"
 	"fmt"
 	"math"
 	"math/rand"
@@ -92,6 +91,7 @@ import (
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/dbterror"
 	"github.com/pingcap/tidb/util/execdetails"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/pingcap/tidb/util/kvcache"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/logutil/consistency"
@@ -4217,7 +4217,7 @@ func (s *session) setRequestSource(ctx context.Context, stmtLabel string, stmtNo
 	}
 	// panic in test mode in case there are requests without source in the future.
 	// log warnings in production mode.
-	if flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil {
+	if intest.InTest {
 		panic("unexpected no source type context, if you see this error, " +
 			"the `RequestSourceTypeKey` is missing in your context")
 	} else {

testkit/BUILD.bazel

@@ -31,6 +31,7 @@ go_library(
         "//util/breakpoint",
         "//util/chunk",
         "//util/gctuner",
+        "//util/intest",
         "//util/mathutil",
         "//util/sqlexec",
         "@com_github_pingcap_errors//:errors",

testkit/testkit.go

@@ -33,6 +33,7 @@ import (
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tidb/util/intest"
 	"github.com/pingcap/tidb/util/mathutil"
 	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/stretchr/testify/assert"
@@ -56,6 +57,7 @@ type TestKit struct {
 
 // NewTestKit returns a new *TestKit.
 func NewTestKit(t testing.TB, store kv.Storage) *TestKit {
+	require.True(t, intest.InTest, "you should add --tags=intest when to test")
 	runtime.GOMAXPROCS(mathutil.Min(16, runtime.GOMAXPROCS(0)))
 	tk := &TestKit{
 		require: require.New(t),

tools/check/ut.go

@@ -837,7 +837,7 @@ func skipDIR(pkg string) bool {
 
 func buildTestBinary(pkg string) error {
 	// go test -c
-	cmd := exec.Command("go", "test", "-c", "-vet", "off", "-o", testFileName(pkg))
+	cmd := exec.Command("go", "test", "-c", "-vet", "off", "--tags=intest", "-o", testFileName(pkg))
 	if coverprofile != "" {
 		cmd.Args = append(cmd.Args, "-cover")
 	}

util/intest/BUILD.bazel

@@ -0,0 +1,11 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "intest",
+    srcs = [
+        "common.go",
+        "intest.go",  #keep
+    ],
+    importpath = "github.com/pingcap/tidb/util/intest",
+    visibility = ["//visibility:public"],
+)

util/intest/common.go

@@ -0,0 +1,20 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !intest
+
+package intest
+
+// InTest checks if the code is running in test.
+const InTest = false

util/intest/intest.go

@@ -0,0 +1,20 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build intest
+
+package intest
+
+// InTest checks if the code is running in test.
+const InTest = true