diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index e4e8d55975a6a..fc4460acb255e 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -566,18 +566,9 @@ func (s *testIntegrationSerialSuite) TestJoinNotSupportedByTiFlash(c *C) {
 	tk.MustExec("analyze table table_1")
 
 	// Create virtual tiflash replica info.
-	dom := domain.GetDomain(tk.Se)
-	is := dom.InfoSchema()
-	db, exists := is.SchemaByName(model.NewCIStr("test"))
-	c.Assert(exists, IsTrue)
-	for _, tblInfo := range db.Tables {
-		if tblInfo.Name.L == "table_1" {
-			tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
-				Count:     1,
-				Available: true,
-			}
-		}
-	}
+	tb := testkit.TestGetTableByName(c, tk.Se, "test", "table_1")
+	err := domain.GetDomain(tk.Se).DDL().UpdateTableReplicaInfo(tk.Se, tb.Meta().ID, true)
+	c.Assert(err, IsNil)
 
 	tk.MustExec("set @@session.tidb_isolation_read_engines = 'tiflash'")
 	tk.MustExec("set @@session.tidb_allow_mpp = 1")
@@ -635,18 +626,9 @@ func (s *testIntegrationSerialSuite) TestMPPNotSupportedInNewCollation(c *C) {
 	tk.MustExec("analyze table table_1")
 
 	// Create virtual tiflash replica info.
-	dom := domain.GetDomain(tk.Se)
-	is := dom.InfoSchema()
-	db, exists := is.SchemaByName(model.NewCIStr("test"))
-	c.Assert(exists, IsTrue)
-	for _, tblInfo := range db.Tables {
-		if tblInfo.Name.L == "table_1" {
-			tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
-				Count:     1,
-				Available: true,
-			}
-		}
-	}
+	tb := testkit.TestGetTableByName(c, tk.Se, "test", "table_1")
+	err := domain.GetDomain(tk.Se).DDL().UpdateTableReplicaInfo(tk.Se, tb.Meta().ID, true)
+	c.Assert(err, IsNil)
 
 	collate.SetNewCollationEnabledForTest(true)
 	tk.MustExec("set @@session.tidb_isolation_read_engines = 'tiflash'")
diff --git a/util/testkit/testkit.go b/util/testkit/testkit.go
index e82ad885fe939..9e9bb21218441 100644
--- a/util/testkit/testkit.go
+++ b/util/testkit/testkit.go
@@ -30,7 +30,9 @@ import (
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/session"
+	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/pingcap/tidb/util/testutil"
@@ -367,3 +369,14 @@ func WithPruneMode(tk *TestKit, mode variable.PartitionPruneMode, f func()) {
 	tk.MustExec("set global tidb_partition_prune_mode=`" + string(mode) + "`")
 	f()
 }
+
+// TestGetTableByName get the table info by name, will assert error if not found
+func TestGetTableByName(c *check.C, ctx sessionctx.Context, db, table string) table.Table {
+	dom := domain.GetDomain(ctx)
+	// Make sure the table schema is the new schema.
+	err := dom.Reload()
+	c.Assert(err, check.IsNil)
+	tbl, err := dom.InfoSchema().TableByName(model.NewCIStr(db), model.NewCIStr(table))
+	c.Assert(err, check.IsNil)
+	return tbl
+}