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

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package restore_test
     4  
     5  import (
     6  	"bytes"
     7  
     8  	. "github.com/pingcap/check"
     9  	"github.com/pingcap/kvproto/pkg/import_sstpb"
    10  	"github.com/pingcap/tidb/tablecodec"
    11  
    12  	"github.com/pingcap/br/pkg/restore"
    13  	"github.com/pingcap/br/pkg/rtree"
    14  )
    15  
    16  type testRangeSuite struct{}
    17  
    18  var _ = Suite(&testRangeSuite{})
    19  
    20  type rangeEquals struct {
    21  	*CheckerInfo
    22  }
    23  
    24  var RangeEquals Checker = &rangeEquals{
    25  	&CheckerInfo{Name: "RangeEquals", Params: []string{"obtained", "expected"}},
    26  }
    27  
    28  func (checker *rangeEquals) Check(params []interface{}, names []string) (result bool, error string) {
    29  	obtained := params[0].([]rtree.Range)
    30  	expected := params[1].([]rtree.Range)
    31  	if len(obtained) != len(expected) {
    32  		return false, ""
    33  	}
    34  	for i := range obtained {
    35  		if !bytes.Equal(obtained[i].StartKey, expected[i].StartKey) ||
    36  			!bytes.Equal(obtained[i].EndKey, expected[i].EndKey) {
    37  			return false, ""
    38  		}
    39  	}
    40  	return true, ""
    41  }
    42  
    43  func (s *testRangeSuite) TestSortRange(c *C) {
    44  	dataRules := []*import_sstpb.RewriteRule{
    45  		{OldKeyPrefix: tablecodec.GenTableRecordPrefix(1), NewKeyPrefix: tablecodec.GenTableRecordPrefix(4)},
    46  		{OldKeyPrefix: tablecodec.GenTableRecordPrefix(2), NewKeyPrefix: tablecodec.GenTableRecordPrefix(5)},
    47  	}
    48  	rewriteRules := &restore.RewriteRules{
    49  		Data: dataRules,
    50  	}
    51  	ranges1 := []rtree.Range{
    52  		{
    53  			StartKey: append(tablecodec.GenTableRecordPrefix(1), []byte("aaa")...),
    54  			EndKey:   append(tablecodec.GenTableRecordPrefix(1), []byte("bbb")...), Files: nil,
    55  		},
    56  	}
    57  	rs1, err := restore.SortRanges(ranges1, rewriteRules)
    58  	c.Assert(err, IsNil, Commentf("sort range1 failed: %v", err))
    59  	c.Assert(rs1, RangeEquals, []rtree.Range{
    60  		{
    61  			StartKey: append(tablecodec.GenTableRecordPrefix(4), []byte("aaa")...),
    62  			EndKey:   append(tablecodec.GenTableRecordPrefix(4), []byte("bbb")...), Files: nil,
    63  		},
    64  	})
    65  
    66  	ranges2 := []rtree.Range{
    67  		{
    68  			StartKey: append(tablecodec.GenTableRecordPrefix(1), []byte("aaa")...),
    69  			EndKey:   append(tablecodec.GenTableRecordPrefix(2), []byte("bbb")...), Files: nil,
    70  		},
    71  	}
    72  	_, err = restore.SortRanges(ranges2, rewriteRules)
    73  	c.Assert(err, ErrorMatches, "table id mismatch.*")
    74  
    75  	ranges3 := initRanges()
    76  	rewriteRules1 := initRewriteRules()
    77  	rs3, err := restore.SortRanges(ranges3, rewriteRules1)
    78  	c.Assert(err, IsNil, Commentf("sort range1 failed: %v", err))
    79  	c.Assert(rs3, RangeEquals, []rtree.Range{
    80  		{StartKey: []byte("bbd"), EndKey: []byte("bbf"), Files: nil},
    81  		{StartKey: []byte("bbf"), EndKey: []byte("bbj"), Files: nil},
    82  		{StartKey: []byte("xxa"), EndKey: []byte("xxe"), Files: nil},
    83  		{StartKey: []byte("xxe"), EndKey: []byte("xxz"), Files: nil},
    84  	})
    85  }