diff --git a/planner/core/logical_plans_test.go b/planner/core/logical_plans_test.go
index c3193d34f659a..c9d86355238dc 100644
--- a/planner/core/logical_plans_test.go
+++ b/planner/core/logical_plans_test.go
@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/testkit/testdata"
 	"github.com/pingcap/tidb/util/hint"
+	"github.com/pingcap/tidb/util/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -91,7 +92,22 @@ func createPlannerSuite() (s *plannerSuite) {
 		}
 	}
 	s.is = infoschema.MockInfoSchema(tblInfos)
-	s.ctx = MockContext()
+	ctx := mock.NewContext()
+	ctx.Store = &mock.Store{
+		Client: &mock.Client{},
+	}
+	initStatsCtx := mock.NewContext()
+	initStatsCtx.Store = &mock.Store{
+		Client: &mock.Client{},
+	}
+	ctx.GetSessionVars().CurrentDB = "test"
+	do := domain.NewMockDomain()
+	if err := do.CreateStatsHandle(ctx, initStatsCtx); err != nil {
+		panic(fmt.Sprintf("create mock context panic: %+v", err))
+	}
+	domain.BindDomain(ctx, do)
+	ctx.SetInfoSchema(s.is)
+	s.ctx = ctx
 	domain.GetDomain(s.ctx).MockInfoCacheAndLoadInfoSchema(s.is)
 	s.ctx.GetSessionVars().EnableWindowFunction = true
 	s.p = parser.New()
diff --git a/planner/core/point_get_plan.go b/planner/core/point_get_plan.go
index 90de4ab838041..1721d07d05d31 100644
--- a/planner/core/point_get_plan.go
+++ b/planner/core/point_get_plan.go
@@ -1656,7 +1656,7 @@ func buildPointUpdatePlan(ctx sessionctx.Context, pointPlan PhysicalPlan, dbName
 		VirtualAssignmentsOffset:  len(orderedList),
 	}.Init(ctx)
 	updatePlan.names = pointPlan.OutputNames()
-	is := ctx.GetInfoSchema().(infoschema.InfoSchema)
+	is := sessiontxn.GetTxnManager(ctx).GetTxnInfoSchema()
 	t, _ := is.TableByID(tbl.ID)
 	updatePlan.tblID2Table = map[int64]table.Table{
 		tbl.ID: t,
diff --git a/util/mock/context.go b/util/mock/context.go
index 698ef33b0ce65..b0d02ca659495 100644
--- a/util/mock/context.go
+++ b/util/mock/context.go
@@ -453,6 +453,11 @@ func (c *Context) InSandBoxMode() bool {
 	return c.inSandBoxMode
 }
 
+// SetInfoSchema is to set info shema for the test.
+func (c *Context) SetInfoSchema(is sessionctx.InfoschemaMetaVersion) {
+	c.is = is
+}
+
 // Close implements the sessionctx.Context interface.
 func (*Context) Close() {}