test: stabilize TestGetTargetNodeCpuCnt (pingcap#57823)

close pingcap#57822
RidRisR · Nov 29, 2024 · b9eeab4 · b9eeab4
1 parent f6ff412
commit b9eeab4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel
@@ -135,6 +135,7 @@ go_test(
         "//pkg/session",
         "//pkg/sessionctx/variable",
         "//pkg/testkit",
+        "//pkg/testkit/testfailpoint",
         "//pkg/testkit/testsetup",
         "//pkg/types",
         "//pkg/util",

diff --git a/pkg/executor/importer/importer_testkit_test.go b/pkg/executor/importer/importer_testkit_test.go
@@ -45,6 +45,7 @@ import (
 	"github.com/pingcap/tidb/pkg/session"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/dbterror/exeerrors"
 	"github.com/pingcap/tidb/pkg/util/logutil"
@@ -159,18 +160,14 @@ func TestVerifyChecksum(t *testing.T) {
 }
 
 func TestGetTargetNodeCpuCnt(t *testing.T) {
-	_, tm, ctx := testutil.InitTableTest(t)
-	old := variable.EnableDistTask.Load()
+	store, tm, ctx := testutil.InitTableTest(t)
+	tk := testkit.NewTestKit(t, store)
 
-	variable.EnableDistTask.Store(false)
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu", "return(16)"))
-	t.Cleanup(func() {
-		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu"))
-		variable.EnableDistTask.Store(old)
-	})
+	tk.MustExec("set @@global.tidb_enable_dist_task = off;")
+	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu", "return(16)")
 	require.NoError(t, tm.InitMeta(ctx, "tidb1", ""))
 
-	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu", "return(8)"))
+	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/util/cpu/mockNumCpu", "return(8)")
 	targetNodeCPUCnt, err := importer.GetTargetNodeCPUCnt(ctx, importer.DataSourceTypeQuery, "")
 	require.NoError(t, err)
 	require.Equal(t, 8, targetNodeCPUCnt)
@@ -187,7 +184,7 @@ func TestGetTargetNodeCpuCnt(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, 8, targetNodeCPUCnt)
 	// disttask enabled
-	variable.EnableDistTask.Store(true)
+	tk.MustExec("set @@global.tidb_enable_dist_task = on;")
 
 	targetNodeCPUCnt, err = importer.GetTargetNodeCPUCnt(ctx, importer.DataSourceTypeFile, "s3://path/to/xxx.csv")
 	require.NoError(t, err)