diff --git a/integration_tests/assertion_test.go b/integration_tests/assertion_test.go
index 8a3b718f..3a2fab0c 100644
--- a/integration_tests/assertion_test.go
+++ b/integration_tests/assertion_test.go
@@ -43,7 +43,7 @@ type testAssertionSuite struct {
 }
 
 func (s *testAssertionSuite) SetupTest() {
-	s.store = tikv.StoreProbe{KVStore: NewTestStore(s.T())}
+	s.store = tikv.StoreProbe{KVStore: NewTestUniStore(s.T())}
 }
 
 func (s *testAssertionSuite) TearDownTest() {
diff --git a/integration_tests/lock_test.go b/integration_tests/lock_test.go
index 2317e6cf..59bcc833 100644
--- a/integration_tests/lock_test.go
+++ b/integration_tests/lock_test.go
@@ -71,7 +71,7 @@ type testLockSuite struct {
 }
 
 func (s *testLockSuite) SetupTest() {
-	s.store = tikv.StoreProbe{KVStore: NewTestStore(s.T())}
+	s.store = tikv.StoreProbe{KVStore: NewTestUniStore(s.T())}
 }
 
 func (s *testLockSuite) TearDownTest() {
diff --git a/integration_tests/prewrite_test.go b/integration_tests/prewrite_test.go
index 102e4e0c..a2f38c98 100644
--- a/integration_tests/prewrite_test.go
+++ b/integration_tests/prewrite_test.go
@@ -38,9 +38,9 @@ import (
 	"testing"
 
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
-	"github.com/pingcap/tidb/store/mockstore/unistore"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/testutils"
 	"github.com/tikv/client-go/v2/tikv"
 	"github.com/tikv/client-go/v2/txnkv/transaction"
 )
@@ -48,9 +48,9 @@ import (
 func TestSetMinCommitTSInAsyncCommit(t *testing.T) {
 	require, assert := require.New(t), assert.New(t)
 
-	client, pdClient, cluster, err := unistore.New("")
+	client, cluster, pdClient, err := testutils.NewMockTiKV("", nil)
 	require.Nil(err)
-	unistore.BootstrapWithSingleStore(cluster)
+	testutils.BootstrapWithSingleStore(cluster)
 	store, err := tikv.NewTestTiKVStore(client, pdClient, nil, nil, 0)
 	require.Nil(err)
 	defer store.Close()
diff --git a/integration_tests/snapshot_fail_test.go b/integration_tests/snapshot_fail_test.go
index fa6745dc..e339d1e2 100644
--- a/integration_tests/snapshot_fail_test.go
+++ b/integration_tests/snapshot_fail_test.go
@@ -42,7 +42,6 @@ import (
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
-	"github.com/pingcap/tidb/store/mockstore/unistore"
 	"github.com/stretchr/testify/suite"
 	tikverr "github.com/tikv/client-go/v2/error"
 	"github.com/tikv/client-go/v2/oracle"
@@ -62,11 +61,7 @@ type testSnapshotFailSuite struct {
 }
 
 func (s *testSnapshotFailSuite) SetupSuite() {
-	client, pdClient, cluster, err := unistore.New("")
-	s.Require().Nil(err)
-	unistore.BootstrapWithSingleStore(cluster)
-	store, err := tikv.NewTestTiKVStore(fpClient{Client: client}, pdClient, nil, nil, 0)
-	s.Require().Nil(err)
+	store := NewTestUniStore(s.T())
 	s.store = tikv.StoreProbe{KVStore: store}
 }
 
diff --git a/integration_tests/util_test.go b/integration_tests/util_test.go
index 4d1046ba..c89656b7 100644
--- a/integration_tests/util_test.go
+++ b/integration_tests/util_test.go
@@ -47,6 +47,7 @@ import (
 	"github.com/pingcap/tidb/store/mockstore/unistore"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/config"
+	"github.com/tikv/client-go/v2/testutils"
 	"github.com/tikv/client-go/v2/tikv"
 	"github.com/tikv/client-go/v2/txnkv/transaction"
 	"github.com/tikv/client-go/v2/util/codec"
@@ -65,19 +66,25 @@ func NewTestStore(t *testing.T) *tikv.KVStore {
 	}
 
 	if *withTiKV {
-		addrs := strings.Split(*pdAddrs, ",")
-		pdClient, err := pd.NewClient(addrs, pd.SecurityOption{})
-		require.Nil(t, err)
-		var securityConfig config.Security
-		tlsConfig, err := securityConfig.ToTLSConfig()
-		require.Nil(t, err)
-		spKV, err := tikv.NewEtcdSafePointKV(addrs, tlsConfig)
-		require.Nil(t, err)
-		store, err := tikv.NewKVStore("test-store", &tikv.CodecPDClient{Client: pdClient}, spKV, tikv.NewRPCClient())
-		require.Nil(t, err)
-		err = clearStorage(store)
-		require.Nil(t, err)
-		return store
+		return newTiKVStore(t)
+	}
+	client, cluster, pdClient, err := testutils.NewMockTiKV("", nil)
+	require.NoError(t, err)
+	testutils.BootstrapWithSingleStore(cluster)
+	store, err := tikv.NewTestTiKVStore(client, pdClient, nil, nil, 0)
+	require.Nil(t, err)
+	return store
+}
+
+// NewTestUniStore creates a KVStore (using tidb/unistore) for testing purpose.
+// TODO: switch to use mockstore and remove it.
+func NewTestUniStore(t *testing.T) *tikv.KVStore {
+	if !flag.Parsed() {
+		flag.Parse()
+	}
+
+	if *withTiKV {
+		return newTiKVStore(t)
 	}
 	client, pdClient, cluster, err := unistore.New("")
 	require.Nil(t, err)
@@ -87,6 +94,22 @@ func NewTestStore(t *testing.T) *tikv.KVStore {
 	return store
 }
 
+func newTiKVStore(t *testing.T) *tikv.KVStore {
+	addrs := strings.Split(*pdAddrs, ",")
+	pdClient, err := pd.NewClient(addrs, pd.SecurityOption{})
+	require.Nil(t, err)
+	var securityConfig config.Security
+	tlsConfig, err := securityConfig.ToTLSConfig()
+	require.Nil(t, err)
+	spKV, err := tikv.NewEtcdSafePointKV(addrs, tlsConfig)
+	require.Nil(t, err)
+	store, err := tikv.NewKVStore("test-store", &tikv.CodecPDClient{Client: pdClient}, spKV, tikv.NewRPCClient())
+	require.Nil(t, err)
+	err = clearStorage(store)
+	require.Nil(t, err)
+	return store
+}
+
 func clearStorage(store *tikv.KVStore) error {
 	txn, err := store.Begin()
 	if err != nil {