github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/conn/conn_test.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package conn
     4  
     5  import (
     6  	"context"
     7  	"testing"
     8  
     9  	"github.com/pingcap/br/pkg/pdutil"
    10  
    11  	. "github.com/pingcap/check"
    12  	"github.com/pingcap/kvproto/pkg/metapb"
    13  	pd "github.com/tikv/pd/client"
    14  )
    15  
    16  func TestT(t *testing.T) {
    17  	TestingT(t)
    18  }
    19  
    20  var _ = Suite(&testClientSuite{})
    21  
    22  type testClientSuite struct {
    23  	ctx    context.Context
    24  	cancel context.CancelFunc
    25  
    26  	mgr *Mgr
    27  }
    28  
    29  func (s *testClientSuite) SetUpSuite(c *C) {
    30  	s.ctx, s.cancel = context.WithCancel(context.Background())
    31  	s.mgr = &Mgr{PdController: &pdutil.PdController{}}
    32  }
    33  
    34  func (s *testClientSuite) TearDownSuite(c *C) {
    35  	s.cancel()
    36  }
    37  
    38  type fakePDClient struct {
    39  	pd.Client
    40  	stores []*metapb.Store
    41  }
    42  
    43  func (fpdc fakePDClient) GetAllStores(context.Context, ...pd.GetStoreOption) ([]*metapb.Store, error) {
    44  	return append([]*metapb.Store{}, fpdc.stores...), nil
    45  }
    46  
    47  func (s *testClientSuite) TestGetAllTiKVStores(c *C) {
    48  	testCases := []struct {
    49  		stores         []*metapb.Store
    50  		storeBehavior  StoreBehavior
    51  		expectedStores map[uint64]int
    52  		expectedError  string
    53  	}{
    54  		{
    55  			stores: []*metapb.Store{
    56  				{Id: 1},
    57  			},
    58  			storeBehavior:  SkipTiFlash,
    59  			expectedStores: map[uint64]int{1: 1},
    60  		},
    61  		{
    62  			stores: []*metapb.Store{
    63  				{Id: 1},
    64  			},
    65  			storeBehavior:  ErrorOnTiFlash,
    66  			expectedStores: map[uint64]int{1: 1},
    67  		},
    68  		{
    69  			stores: []*metapb.Store{
    70  				{Id: 1},
    71  				{Id: 2, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}},
    72  			},
    73  			storeBehavior:  SkipTiFlash,
    74  			expectedStores: map[uint64]int{1: 1},
    75  		},
    76  		{
    77  			stores: []*metapb.Store{
    78  				{Id: 1},
    79  				{Id: 2, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}},
    80  			},
    81  			storeBehavior: ErrorOnTiFlash,
    82  			expectedError: "cannot restore to a cluster with active TiFlash stores.*",
    83  		},
    84  		{
    85  			stores: []*metapb.Store{
    86  				{Id: 1},
    87  				{Id: 2, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}},
    88  				{Id: 3},
    89  				{Id: 4, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tikv"}}},
    90  				{Id: 5, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tikv"}, {Key: "engine", Value: "tiflash"}}},
    91  				{Id: 6, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tiflash"}, {Key: "engine", Value: "tikv"}}},
    92  			},
    93  			storeBehavior:  SkipTiFlash,
    94  			expectedStores: map[uint64]int{1: 1, 3: 1, 4: 1, 6: 1},
    95  		},
    96  		{
    97  			stores: []*metapb.Store{
    98  				{Id: 1},
    99  				{Id: 2, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}},
   100  				{Id: 3},
   101  				{Id: 4, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tikv"}}},
   102  				{Id: 5, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tikv"}, {Key: "engine", Value: "tiflash"}}},
   103  				{Id: 6, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tiflash"}, {Key: "engine", Value: "tikv"}}},
   104  			},
   105  			storeBehavior: ErrorOnTiFlash,
   106  			expectedError: "cannot restore to a cluster with active TiFlash stores.*",
   107  		},
   108  		{
   109  			stores: []*metapb.Store{
   110  				{Id: 1},
   111  				{Id: 2, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}},
   112  				{Id: 3},
   113  				{Id: 4, Labels: []*metapb.StoreLabel{{Key: "engine", Value: "tikv"}}},
   114  				{Id: 5, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tikv"}, {Key: "engine", Value: "tiflash"}}},
   115  				{Id: 6, Labels: []*metapb.StoreLabel{{Key: "else", Value: "tiflash"}, {Key: "engine", Value: "tikv"}}},
   116  			},
   117  			storeBehavior:  TiFlashOnly,
   118  			expectedStores: map[uint64]int{2: 1, 5: 1},
   119  		},
   120  	}
   121  
   122  	for _, testCase := range testCases {
   123  		pdClient := fakePDClient{stores: testCase.stores}
   124  		stores, err := GetAllTiKVStores(context.Background(), pdClient, testCase.storeBehavior)
   125  		if len(testCase.expectedError) != 0 {
   126  			c.Assert(err, ErrorMatches, testCase.expectedError)
   127  			continue
   128  		}
   129  		foundStores := make(map[uint64]int)
   130  		for _, store := range stores {
   131  			foundStores[store.Id]++
   132  		}
   133  		c.Assert(foundStores, DeepEquals, testCase.expectedStores)
   134  	}
   135  }
   136  
   137  func (s *testClientSuite) TestGetConnOnCanceledContext(c *C) {
   138  	ctx, cancel := context.WithCancel(context.Background())
   139  	cancel()
   140  
   141  	_, err := s.mgr.GetBackupClient(ctx, 42)
   142  	c.Assert(err, ErrorMatches, ".*context canceled.*")
   143  	_, err = s.mgr.ResetBackupClient(ctx, 42)
   144  	c.Assert(err, ErrorMatches, ".*context canceled.*")
   145  }