diff --git a/pkg/planner/core/logical_plans_test.go b/pkg/planner/core/logical_plans_test.go
index 82283419f6012..770b433d3a38c 100644
--- a/pkg/planner/core/logical_plans_test.go
+++ b/pkg/planner/core/logical_plans_test.go
@@ -37,6 +37,7 @@ import (
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/testkit/testdata"
 	"github.com/pingcap/tidb/pkg/util/hint"
+	"github.com/pingcap/tidb/pkg/util/mock"
 	"github.com/pingcap/tipb/go-tipb"
 	"github.com/stretchr/testify/require"
 )
@@ -93,7 +94,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/pkg/planner/core/point_get_plan.go b/pkg/planner/core/point_get_plan.go
index 63403cfe5256f..0fe6a7bfbe83a 100644
--- a/pkg/planner/core/point_get_plan.go
+++ b/pkg/planner/core/point_get_plan.go
@@ -1657,7 +1657,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/pkg/util/mock/context.go b/pkg/util/mock/context.go
index 56dc4ae1ad50d..24bf504cb1fa6 100644
--- a/pkg/util/mock/context.go
+++ b/pkg/util/mock/context.go
@@ -450,6 +450,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() {}