mock: add more mock to avoid dm ci failed (pingcap#50939)

close pingcap#50938
jiyfhust · Feb 17, 2024 · ba9b8a8 · ba9b8a8
1 parent 17a25f4
commit ba9b8a8
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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")