store/tikv: remove dependency on testkit (#22520)

store/tikv/batch_coprocessor_test.go → store/batch_coprocessor_test.go

@@ -11,7 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tikv_test
+package store
 
 import (
 	"context"

store/tikv/sql_fail_test.go → store/sql_fail_test.go

@@ -11,7 +11,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tikv_test
+package store
 
 import (
 	"context"
@@ -26,7 +26,7 @@ import (
 	"github.com/pingcap/parser/terror"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/session"
-	. "github.com/pingcap/tidb/store/tikv"
+	"github.com/pingcap/tidb/store/tikv"
 	"github.com/pingcap/tidb/util/mock"
 	"github.com/pingcap/tidb/util/testkit"
 )
@@ -44,14 +44,14 @@ type testSQLSerialSuite struct {
 
 type testSQLSuiteBase struct {
 	OneByOneSuite
-	store Storage
+	store tikv.Storage
 	dom   *domain.Domain
 }
 
 func (s *testSQLSuiteBase) SetUpSuite(c *C) {
 	s.OneByOneSuite.SetUpSuite(c)
 	var err error
-	s.store = NewTestStore(c).(Storage)
+	s.store = NewTestStore(c).(tikv.Storage)
 	// actual this is better done in `OneByOneSuite.SetUpSuite`, but this would cause circle dependency
 	if *WithTiKV {
 		session.ResetStoreForWithTiKVTest(s.store)
@@ -99,7 +99,7 @@ func (s *testSQLSerialSuite) TestFailBusyServerCop(c *C) {
 }
 
 func TestMain(m *testing.M) {
-	ReadTimeoutMedium = 2 * time.Second
+	tikv.ReadTimeoutMedium = 2 * time.Second
 	os.Exit(m.Run())
 }
 
@@ -116,7 +116,7 @@ func (s *testSQLSuite) TestCoprocessorStreamRecvTimeout(c *C) {
 		enable := true
 		visited := make(chan int, 1)
 		timeouted := false
-		timeout := ReadTimeoutMedium + 100*time.Second
+		timeout := tikv.ReadTimeoutMedium + 100*time.Second
 		ctx := context.WithValue(context.Background(), mock.HookKeyForTest("mockTiKVStreamRecvHook"), func(ctx context.Context) {
 			if !enable {
 				return

store/util_test.go

@@ -0,0 +1,92 @@
+// 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 store
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"sync"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/store/mockstore/unistore"
+	"github.com/pingcap/tidb/store/tikv"
+)
+
+var (
+	withTiKVGlobalLock sync.RWMutex
+	WithTiKV           = flag.Bool("with-tikv", false, "run tests with TiKV cluster started. (not use the mock server)")
+	pdAddrs            = flag.String("pd-addrs", "127.0.0.1:2379", "pd addrs")
+)
+
+// NewTestStore creates a kv.Storage for testing purpose.
+func NewTestStore(c *C) kv.Storage {
+	if !flag.Parsed() {
+		flag.Parse()
+	}
+
+	if *WithTiKV {
+		var d tikv.Driver
+		store, err := d.Open(fmt.Sprintf("tikv://%s", *pdAddrs))
+		c.Assert(err, IsNil)
+		err = clearStorage(store)
+		c.Assert(err, IsNil)
+		return store
+	}
+	client, pdClient, cluster, err := unistore.New("")
+	c.Assert(err, IsNil)
+	unistore.BootstrapWithSingleStore(cluster)
+	store, err := tikv.NewTestTiKVStore(client, pdClient, nil, nil, 0)
+	c.Assert(err, IsNil)
+	return store
+}
+
+func clearStorage(store kv.Storage) error {
+	txn, err := store.Begin()
+	if err != nil {
+		return errors.Trace(err)
+	}
+	iter, err := txn.Iter(nil, nil)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	for iter.Valid() {
+		txn.Delete(iter.Key())
+		if err := iter.Next(); err != nil {
+			return errors.Trace(err)
+		}
+	}
+	return txn.Commit(context.Background())
+}
+
+// OneByOneSuite is a suite, When with-tikv flag is true, there is only one storage, so the test suite have to run one by one.
+type OneByOneSuite struct{}
+
+func (s *OneByOneSuite) SetUpSuite(c *C) {
+	if *WithTiKV {
+		withTiKVGlobalLock.Lock()
+	} else {
+		withTiKVGlobalLock.RLock()
+	}
+}
+
+func (s *OneByOneSuite) TearDownSuite(c *C) {
+	if *WithTiKV {
+		withTiKVGlobalLock.Unlock()
+	} else {
+		withTiKVGlobalLock.RUnlock()
+	}
+}