Fix the tests

Signed-off-by: JmPotato <[email protected]>
pingcap · Nov 14, 2023 · 57d8560 · 57d8560
1 parent 77c52e4
commit 57d8560
Show file tree

Hide file tree

Showing 10 changed files with 52 additions and 32 deletions.
diff --git a/pkg/executor/infoschema_cluster_table_test.go b/pkg/executor/infoschema_cluster_table_test.go
@@ -34,10 +34,12 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/server"
 	"github.com/pingcap/tidb/pkg/store/helper"
+	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/pdapi"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/tikv"
 	pd "github.com/tikv/pd/client/http"
 	"google.golang.org/grpc"
 )
@@ -54,9 +56,12 @@ type infosSchemaClusterTableSuite struct {
 
 func createInfosSchemaClusterTableSuite(t *testing.T) *infosSchemaClusterTableSuite {
 	s := new(infosSchemaClusterTableSuite)
-	s.store, s.dom = testkit.CreateMockStoreAndDomain(t)
-	s.rpcServer, s.listenAddr = setUpRPCService(t, s.dom, "127.0.0.1:0")
 	s.httpServer, s.mockAddr = s.setUpMockPDHTTPServer()
+	s.store, s.dom = testkit.CreateMockStoreAndDomain(
+		t,
+		mockstore.WithTiKVOptions(tikv.WithPDHTTPClient([]string{s.mockAddr})),
+	)
+	s.rpcServer, s.listenAddr = setUpRPCService(t, s.dom, "127.0.0.1:0")
 	s.startTime = time.Now()
 	t.Cleanup(func() {
 		if s.rpcServer != nil {

diff --git a/pkg/executor/tikv_regions_peers_table_test.go b/pkg/executor/tikv_regions_peers_table_test.go
@@ -26,9 +26,11 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/pingcap/fn"
 	"github.com/pingcap/tidb/pkg/store/helper"
+	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/util/pdapi"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/tikv"
 	pd "github.com/tikv/pd/client/http"
 )
 
@@ -110,7 +112,8 @@ func TestTikvRegionPeers(t *testing.T) {
 	router.HandleFunc(pdapi.RegionByID+"/"+"{id}", regionsInfoHandler)
 	defer server.Close()
 
-	store := testkit.CreateMockStore(t)
+	store := testkit.CreateMockStore(t,
+		mockstore.WithTiKVOptions(tikv.WithPDHTTPClient([]string{mockAddr})))
 
 	store = &mockStore{
 		store.(helper.Storage),

diff --git a/pkg/server/handler/tests/BUILD.bazel b/pkg/server/handler/tests/BUILD.bazel
@@ -28,7 +28,6 @@ go_test(
         "//pkg/server/handler/optimizor",
         "//pkg/server/handler/tikvhandler",
         "//pkg/server/internal/testserverclient",
-        "//pkg/server/internal/testutil",
         "//pkg/server/internal/util",
         "//pkg/session",
         "//pkg/sessionctx",

diff --git a/pkg/server/handler/tests/http_handler_test.go b/pkg/server/handler/tests/http_handler_test.go
@@ -50,7 +50,6 @@ import (
 	"github.com/pingcap/tidb/pkg/server/handler/optimizor"
 	"github.com/pingcap/tidb/pkg/server/handler/tikvhandler"
 	"github.com/pingcap/tidb/pkg/server/internal/testserverclient"
-	"github.com/pingcap/tidb/pkg/server/internal/testutil"
 	"github.com/pingcap/tidb/pkg/server/internal/util"
 	"github.com/pingcap/tidb/pkg/session"
 	"github.com/pingcap/tidb/pkg/sessionctx"
@@ -450,22 +449,26 @@ func TestBinlogRecover(t *testing.T) {
 
 func (ts *basicHTTPHandlerTestSuite) startServer(t *testing.T) {
 	var err error
-	ts.store, err = mockstore.NewMockStore()
+	ts.Port = uint(rand.Int31n(50000)) + 10000
+	ts.StatusPort = ts.Port + 1
+	ts.store, err = mockstore.NewMockStore(
+		mockstore.WithTiKVOptions(
+			tikv.WithPDHTTPClient([]string{ts.Addr()}),
+		),
+	)
 	require.NoError(t, err)
 	ts.domain, err = session.BootstrapSession(ts.store)
 	require.NoError(t, err)
 	ts.tidbdrv = server2.NewTiDBDriver(ts.store)
 
 	cfg := util.NewTestConfig()
 	cfg.Store = "tikv"
-	cfg.Port = 0
-	cfg.Status.StatusPort = 0
+	cfg.Port = ts.Port
+	cfg.Status.StatusPort = ts.StatusPort
 	cfg.Status.ReportStatus = true
 
 	server, err := server2.NewServer(cfg, ts.tidbdrv)
 	require.NoError(t, err)
-	ts.Port = testutil.GetPortFromTCPAddr(server.ListenAddr())
-	ts.StatusPort = testutil.GetPortFromTCPAddr(server.StatusListenerAddr())
 	ts.server = server
 	ts.server.SetDomain(ts.domain)
 	go func() {

diff --git a/pkg/server/internal/testserverclient/server_client.go b/pkg/server/internal/testserverclient/server_client.go
@@ -72,7 +72,12 @@ func NewTestServerClient() *TestServerClient {
 	}
 }
 
-// statusURL return the full URL of a status path
+// Addr returns the address of the server.
+func (cli *TestServerClient) Addr() string {
+	return fmt.Sprintf("%s://localhost:%d", cli.StatusScheme, cli.Port)
+}
+
+// StatusURL returns the full URL of a status path
 func (cli *TestServerClient) StatusURL(path string) string {
 	return fmt.Sprintf("%s://localhost:%d%s", cli.StatusScheme, cli.StatusPort, path)
 }

diff --git a/pkg/store/helper/BUILD.bazel b/pkg/store/helper/BUILD.bazel
@@ -47,6 +47,7 @@ go_test(
         "@com_github_pingcap_log//:log",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//testutils",
+        "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_pd_client//http",
         "@io_opencensus_go//stats/view",
         "@org_uber_go_goleak//:goleak",

diff --git a/pkg/store/helper/helper_test.go b/pkg/store/helper/helper_test.go
@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/pdapi"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/testutils"
+	"github.com/tikv/client-go/v2/tikv"
 	pd "github.com/tikv/pd/client/http"
 	"go.opencensus.io/stats/view"
 	"go.uber.org/zap"
@@ -116,16 +117,7 @@ func TestTiKVStoresStat(t *testing.T) {
 
 type mockStore struct {
 	helper.Storage
-	pdAddrs      []string
-	pdHTTPClient pd.Client
-}
-
-func newMockStore(storage helper.Storage, pdAddrs []string) *mockStore {
-	return &mockStore{
-		Storage:      storage,
-		pdAddrs:      pdAddrs,
-		pdHTTPClient: pd.NewClient(pdAddrs),
-	}
+	pdAddrs []string
 }
 
 func (s *mockStore) EtcdAddrs() ([]string, error) {
@@ -148,24 +140,22 @@ func (s *mockStore) Describe() string {
 	return ""
 }
 
-func (s *mockStore) GetPDHTTPClient() pd.Client {
-	return s.pdHTTPClient
-}
-
 func createMockStore(t *testing.T) (store helper.Storage) {
+	server := mockPDHTTPServer()
+
+	pdAddrs := []string{"invalid_pd_address", server.URL[len("http://"):]}
 	s, err := mockstore.NewMockStore(
 		mockstore.WithClusterInspector(func(c testutils.Cluster) {
 			mockstore.BootstrapWithMultiRegions(c, []byte("x"))
 		}),
+		mockstore.WithTiKVOptions(tikv.WithPDHTTPClient(pdAddrs)),
 	)
 	require.NoError(t, err)
 
-	server := mockPDHTTPServer()
-
-	store = newMockStore(
+	store = &mockStore{
 		s.(helper.Storage),
-		[]string{"invalid_pd_address", server.URL[len("http://"):]},
-	)
+		pdAddrs,
+	}
 
 	t.Cleanup(func() {
 		server.Close()

diff --git a/pkg/store/mockstore/mockstore.go b/pkg/store/mockstore/mockstore.go
@@ -92,6 +92,7 @@ type mockOptions struct {
 	txnLocalLatches  uint
 	storeType        StoreType
 	ddlCheckerHijack bool
+	tikvOptions      []tikv.Option
 }
 
 // MockTiKVStoreOption is used to control some behavior of mock tikv.
@@ -106,6 +107,13 @@ func WithMultipleOptions(opts ...MockTiKVStoreOption) MockTiKVStoreOption {
 	}
 }
 
+// WithTiKVOptions sets KV options.
+func WithTiKVOptions(opts ...tikv.Option) MockTiKVStoreOption {
+	return func(args *mockOptions) {
+		args.tikvOptions = opts
+	}
+}
+
 // WithClientHijacker hijacks KV client's behavior, makes it easy to simulate the network
 // problem between TiDB and TiKV.
 func WithClientHijacker(hijacker func(tikv.Client) tikv.Client) MockTiKVStoreOption {

diff --git a/pkg/store/mockstore/tikv.go b/pkg/store/mockstore/tikv.go
@@ -32,7 +32,10 @@ func newMockTikvStore(opt *mockOptions) (kv.Storage, error) {
 	}
 	opt.clusterInspector(cluster)
 
-	kvstore, err := tikv.NewTestTiKVStore(newClientRedirector(client), pdClient, opt.clientHijacker, opt.pdClientHijacker, opt.txnLocalLatches)
+	kvstore, err := tikv.NewTestTiKVStore(
+		newClientRedirector(client), pdClient,
+		opt.clientHijacker, opt.pdClientHijacker,
+		opt.txnLocalLatches, opt.tikvOptions...)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/store/mockstore/unistore.go b/pkg/store/mockstore/unistore.go
@@ -33,7 +33,10 @@ func newUnistore(opts *mockOptions) (kv.Storage, error) {
 		Client: pdClient,
 	}
 
-	kvstore, err := tikv.NewTestTiKVStore(newClientRedirector(client), pdClient, opts.clientHijacker, opts.pdClientHijacker, opts.txnLocalLatches)
+	kvstore, err := tikv.NewTestTiKVStore(
+		newClientRedirector(client), pdClient,
+		opts.clientHijacker, opts.pdClientHijacker,
+		opts.txnLocalLatches, opts.tikvOptions...)
 	if err != nil {
 		return nil, err
 	}