From ba9b8a8ed03134ade5fefc7edbdda8f60b71fb04 Mon Sep 17 00:00:00 2001
From: okJiang <819421878@qq.com>
Date: Sat, 17 Feb 2024 19:57:01 +0800
Subject: [PATCH] mock: add more mock to avoid dm ci failed (#50939)

close pingcap/tidb#50938
---
 pkg/ddl/schematracker/BUILD.bazel        |  2 --
 pkg/ddl/schematracker/dm_tracker_test.go |  9 ++-------
 pkg/util/mock/context.go                 | 17 +++++++++++++++++
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/pkg/ddl/schematracker/BUILD.bazel b/pkg/ddl/schematracker/BUILD.bazel
index 5b30ff5659a22..3df075c1c3a17 100644
--- a/pkg/ddl/schematracker/BUILD.bazel
+++ b/pkg/ddl/schematracker/BUILD.bazel
@@ -53,8 +53,6 @@ go_test(
         "//pkg/parser/ast",
         "//pkg/parser/model",
         "//pkg/sessionctx",
-        "//pkg/sessiontxn",
-        "//pkg/testkit",
         "//pkg/util/chunk",
         "//pkg/util/mock",
         "//pkg/util/sqlexec",
diff --git a/pkg/ddl/schematracker/dm_tracker_test.go b/pkg/ddl/schematracker/dm_tracker_test.go
index 352475c4d295a..dc5eb7de269a7 100644
--- a/pkg/ddl/schematracker/dm_tracker_test.go
+++ b/pkg/ddl/schematracker/dm_tracker_test.go
@@ -32,8 +32,6 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/sessionctx"
-	"github.com/pingcap/tidb/pkg/sessiontxn"
-	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/mock"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
@@ -496,15 +494,12 @@ func TestModifyFromNullToNotNull(t *testing.T) {
 
 	sql = "alter table test.t modify column a int not null;"
 	ctx := context.Background()
-	store := testkit.CreateMockStore(t)
-	sess := testkit.NewTestKit(t, store).Session()
-	err := sessiontxn.NewTxn(context.Background(), sess)
-	require.NoError(t, err)
+	sctx := mock.NewContext()
 	p := parser.New()
 	stmt, err := p.ParseOneStmt(sql, "", "")
 	require.NoError(t, err)
 	// converting from NULL to NOT NULL needs to check data, so caller should provide a RestrictedSQLExecutor
-	executorCtx := mockRestrictedSQLExecutor{sess}
+	executorCtx := mockRestrictedSQLExecutor{sctx}
 	err = tracker.AlterTable(ctx, executorCtx, stmt.(*ast.AlterTableStmt))
 	require.NoError(t, err)
 
diff --git a/pkg/util/mock/context.go b/pkg/util/mock/context.go
index d25125520d350..994891278764e 100644
--- a/pkg/util/mock/context.go
+++ b/pkg/util/mock/context.go
@@ -112,6 +112,23 @@ func (txn *wrapTxn) GetTableInfo(id int64) *model.TableInfo {
 	return txn.Transaction.GetTableInfo(id)
 }
 
+// SetDiskFullOpt implements the interface.
+func (*wrapTxn) SetDiskFullOpt(_ kvrpcpb.DiskFullOpt) {}
+
+// SetOption implements the interface.
+func (*wrapTxn) SetOption(_ int, _ any) {}
+
+// StartTS implements the interface.
+func (*wrapTxn) StartTS() uint64 { return uint64(time.Now().UnixNano()) }
+
+// Get implements the interface.
+func (txn *wrapTxn) Get(ctx context.Context, k kv.Key) ([]byte, error) {
+	if txn.Transaction == nil {
+		return nil, nil
+	}
+	return txn.Transaction.Get(ctx, k)
+}
+
 // Execute implements sqlexec.SQLExecutor Execute interface.
 func (*Context) Execute(_ context.Context, _ string) ([]sqlexec.RecordSet, error) {
 	return nil, errors.Errorf("Not Supported")