From b9eeab4accd00bd74165b9faae551574c93dba15 Mon Sep 17 00:00:00 2001 From: tangenta Date: Fri, 29 Nov 2024 15:18:32 +0800 Subject: [PATCH] test: stabilize TestGetTargetNodeCpuCnt (#57823) close pingcap/tidb#57822 --- pkg/executor/importer/BUILD.bazel | 1 + pkg/executor/importer/importer_testkit_test.go | 17 +++++++---------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index 45affc02d7974..018c6ba46f25e 100644 --- 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 index 12c395d4d8ec9..fe99739c27b3b 100644 --- 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)