testkit: make ut works in nextgen after introduce SYSTEM ks

pkg/disttask/importinto/job_testkit_test.go

@@ -46,7 +46,7 @@ func TestSubmitTaskNextgen(t *testing.T) {
 	}
 	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/domain/MockDisableDistTask", "return(true)")
 	require.NoError(t, kvstore.Register(config.StoreTypeUniStore, mockstore.EmbedUnistoreDriver{}))
-	sysKSStore, _ := testkit.CreateNextgenMockStoreAndDomain(t, keyspace.System)
+	sysKSStore, _ := testkit.CreateMockStoreAndDomainForKS(t, keyspace.System)
 	sysKSTK := testkit.NewTestKit(t, sysKSStore)
 	// in uni-store, Store instances are completely isolated, even they have the
 	// same keyspace name, so we store them here and mock the GetStore
@@ -60,7 +60,7 @@ func TestSubmitTaskNextgen(t *testing.T) {
 			}
 		},
 	)
-	userKSStore, _ := testkit.CreateNextgenMockStoreAndDomain(t, "ks")
+	userKSStore, _ := testkit.CreateMockStoreAndDomainForKS(t, "ks")
 	storeMap["ks"] = userKSStore
 	userKSTK := testkit.NewTestKit(t, userKSStore)
 

pkg/domain/crossks/cross_ks_test.go

@@ -49,7 +49,7 @@ func TestManager(t *testing.T) {
 		t.Skip("cross keyspace is not supported in classic kernel")
 	}
 	require.NoError(t, kvstore.Register(config.StoreTypeUniStore, mockstore.EmbedUnistoreDriver{}))
-	sysKSStore, sysKSDom := testkit.CreateNextgenMockStoreAndDomain(t, keyspace.System)
+	sysKSStore, sysKSDom := testkit.CreateMockStoreAndDomainForKS(t, keyspace.System)
 
 	t.Run("same keyspace access", func(t *testing.T) {
 		_, err := sysKSDom.GetKSSessPool(keyspace.System)
@@ -83,7 +83,7 @@ func TestManager(t *testing.T) {
 		)
 
 		for _, ks := range []string{"ks1", "ks2", "ks3"} {
-			userKSStore, _ := testkit.CreateNextgenMockStoreAndDomain(t, ks)
+			userKSStore, _ := testkit.CreateMockStoreAndDomainForKS(t, ks)
 			storeMap[ks] = userKSStore
 
 			// must switch back to SYSTEM keyspace before creating a cross keyspace session

pkg/testkit/mockstore.go

@@ -266,6 +266,9 @@ func NewDistExecutionContextWithLease(t testing.TB, serverNum int, lease time.Du
 
 // CreateMockStoreAndDomain return a new mock kv.Storage and *domain.Domain.
 func CreateMockStoreAndDomain(t testing.TB, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
+	if kerneltype.IsNextGen() {
+		updateConfigForNextgen(t)
+	}
 	store, err := mockstore.NewMockStore(opts...)
 	require.NoError(t, err)
 	dom := bootstrap(t, store, 500*time.Millisecond)
@@ -282,10 +285,12 @@ func CreateMockStoreAndDomain(t testing.TB, opts ...mockstore.MockTiKVStoreOptio
 
 var keyspaceIDAlloc atomic.Int32
 
-// CreateNextgenMockStoreAndDomain return a new mock kv.Storage and *domain.Domain
-// for nextgen.
-func CreateNextgenMockStoreAndDomain(t testing.TB, ks string, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
-	intest.Assert(kerneltype.IsNextGen(), "CreateNextgenMockStoreAndDomain should only be used in nextgen kernel tests")
+// CreateMockStoreAndDomainForKS return a new mock kv.Storage and *domain.Domain
+// some keyspace.
+// this function is mainly for cross keyspace test, so normally you should use
+// CreateMockStoreAndDomain instead,
+func CreateMockStoreAndDomainForKS(t testing.TB, ks string, opts ...mockstore.MockTiKVStoreOption) (kv.Storage, *domain.Domain) {
+	intest.Assert(kerneltype.IsNextGen(), "CreateMockStoreAndDomainForKS should only be used in nextgen kernel tests")
 
 	bak := *config.GetGlobalConfig()
 	t.Cleanup(func() {
@@ -349,6 +354,9 @@ func bootstrap(t testing.TB, store kv.Storage, lease time.Duration) *domain.Doma
 
 // CreateMockStoreWithSchemaLease return a new mock kv.Storage.
 func CreateMockStoreWithSchemaLease(t testing.TB, lease time.Duration, opts ...mockstore.MockTiKVStoreOption) kv.Storage {
+	if kerneltype.IsNextGen() {
+		updateConfigForNextgen(t)
+	}
 	store, _ := CreateMockStoreAndDomainWithSchemaLease(t, lease, opts...)
 	return schematracker.UnwrapStorage(store)
 }
@@ -373,3 +381,17 @@ func SetTiFlashReplica(t testing.TB, dom *domain.Domain, dbName, tableName strin
 		Available: true,
 	}
 }
+
+func updateConfigForNextgen(t testing.TB) {
+	t.Helper()
+	// in nextgen, SYSTEM ks must be bootstrapped first, to make UT easier, we
+	// always run them inside SYSTEM keyspace, if your test requires bootstrapping
+	// multiple keyspace, you should use CreateMockStoreAndDomainForKS instead.
+	bak := *config.GetGlobalConfig()
+	t.Cleanup(func() {
+		config.StoreGlobalConfig(&bak)
+	})
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.KeyspaceName = keyspace.System
+	})
+}