From fc3f8d51b70ea87ee7e5cb695a033d28c3c56cb0 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Fri, 16 Apr 2021 12:05:51 +0800
Subject: [PATCH] mocktikv: extract test utils (#24066)

---
 store/mockstore/mocktikv/executor_test.go  |  4 +-
 store/mockstore/mocktikv/mock_tikv_test.go | 32 +------------
 store/mockstore/mocktikv/utils.go          | 53 ++++++++++++++++++++++
 3 files changed, 56 insertions(+), 33 deletions(-)
 create mode 100644 store/mockstore/mocktikv/utils.go

diff --git a/store/mockstore/mocktikv/executor_test.go b/store/mockstore/mocktikv/executor_test.go
index 83545c2591..a7f17ca8cc 100644
--- a/store/mockstore/mocktikv/executor_test.go
+++ b/store/mockstore/mocktikv/executor_test.go
@@ -85,8 +85,8 @@ func (s *testExecutorSuite) TestResolvedLargeTxnLocks(c *C) {
 
 	// Simulate a large txn (holding a pk lock with large TTL).
 	// Secondary lock 200ms, primary lock 100s
-	mocktikv.MustPrewriteOK(c, s.mvccStore, mocktikv.PutMutations("primary", "value"), "primary", tso, 100000)
-	mocktikv.MustPrewriteOK(c, s.mvccStore, mocktikv.PutMutations(string(key), "value"), "primary", tso, 200)
+	c.Assert(mocktikv.MustPrewrite(s.mvccStore, mocktikv.PutMutations("primary", "value"), "primary", tso, 100000), IsTrue)
+	c.Assert(mocktikv.MustPrewrite(s.mvccStore, mocktikv.PutMutations(string(key), "value"), "primary", tso, 200), IsTrue)
 
 	// Simulate the action of reading meet the lock of a large txn.
 	// The lock of the large transaction should not block read.
diff --git a/store/mockstore/mocktikv/mock_tikv_test.go b/store/mockstore/mocktikv/mock_tikv_test.go
index 7ddcf8f360..ba8a031377 100644
--- a/store/mockstore/mocktikv/mock_tikv_test.go
+++ b/store/mockstore/mocktikv/mock_tikv_test.go
@@ -51,21 +51,6 @@ func (s *testMockTiKVSuite) SetUpTest(c *C) {
 	c.Assert(err, IsNil)
 }
 
-// PutMutations is exported for testing.
-var PutMutations = putMutations
-
-func putMutations(kvpairs ...string) []*kvrpcpb.Mutation {
-	var mutations []*kvrpcpb.Mutation
-	for i := 0; i < len(kvpairs); i += 2 {
-		mutations = append(mutations, &kvrpcpb.Mutation{
-			Op:    kvrpcpb.Op_Put,
-			Key:   []byte(kvpairs[i]),
-			Value: []byte(kvpairs[i+1]),
-		})
-	}
-	return mutations
-}
-
 func lock(key, primary string, ts uint64) *kvrpcpb.LockInfo {
 	return &kvrpcpb.LockInfo{
 		Key:         []byte(key),
@@ -160,27 +145,12 @@ func (s *testMockTiKVSuite) mustRangeReverseScanOK(c *C, start, end string, limi
 	}
 }
 
-func MustPrewriteOK(c *C, store MVCCStore, mutations []*kvrpcpb.Mutation, primary string, startTS uint64, ttl uint64) {
-	s := testMockTiKVSuite{store}
-	s.mustPrewriteWithTTLOK(c, mutations, primary, startTS, ttl)
-}
-
 func (s *testMockTiKVSuite) mustPrewriteOK(c *C, mutations []*kvrpcpb.Mutation, primary string, startTS uint64) {
 	s.mustPrewriteWithTTLOK(c, mutations, primary, startTS, 0)
 }
 
 func (s *testMockTiKVSuite) mustPrewriteWithTTLOK(c *C, mutations []*kvrpcpb.Mutation, primary string, startTS uint64, ttl uint64) {
-	req := &kvrpcpb.PrewriteRequest{
-		Mutations:    mutations,
-		PrimaryLock:  []byte(primary),
-		StartVersion: startTS,
-		LockTtl:      ttl,
-		MinCommitTs:  startTS + 1,
-	}
-	errs := s.store.Prewrite(req)
-	for _, err := range errs {
-		c.Assert(err, IsNil)
-	}
+	c.Assert(mustPrewriteWithTTL(s.store, mutations, primary, startTS, ttl), IsTrue)
 }
 
 func (s *testMockTiKVSuite) mustCommitOK(c *C, keys [][]byte, startTS, commitTS uint64) {
diff --git a/store/mockstore/mocktikv/utils.go b/store/mockstore/mocktikv/utils.go
new file mode 100644
index 0000000000..fbd7c2b818
--- /dev/null
+++ b/store/mockstore/mocktikv/utils.go
@@ -0,0 +1,53 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mocktikv
+
+import "github.com/pingcap/kvproto/pkg/kvrpcpb"
+
+// PutMutations is exported for testing.
+var PutMutations = putMutations
+
+func putMutations(kvpairs ...string) []*kvrpcpb.Mutation {
+	var mutations []*kvrpcpb.Mutation
+	for i := 0; i < len(kvpairs); i += 2 {
+		mutations = append(mutations, &kvrpcpb.Mutation{
+			Op:    kvrpcpb.Op_Put,
+			Key:   []byte(kvpairs[i]),
+			Value: []byte(kvpairs[i+1]),
+		})
+	}
+	return mutations
+}
+
+// MustPrewrite write mutations to mvcc store.
+func MustPrewrite(store MVCCStore, mutations []*kvrpcpb.Mutation, primary string, startTS uint64, ttl uint64) bool {
+	return mustPrewriteWithTTL(store, mutations, primary, startTS, ttl)
+}
+
+func mustPrewriteWithTTL(store MVCCStore, mutations []*kvrpcpb.Mutation, primary string, startTS uint64, ttl uint64) bool {
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    mutations,
+		PrimaryLock:  []byte(primary),
+		StartVersion: startTS,
+		LockTtl:      ttl,
+		MinCommitTs:  startTS + 1,
+	}
+	errs := store.Prewrite(req)
+	for _, err := range errs {
+		if err != nil {
+			return false
+		}
+	}
+	return true
+}