github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/pingcap/tidb/store/tikv/ticlient_test.go (about)

     1  // Copyright 2016 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package tikv
    15  
    16  import (
    17  	"flag"
    18  	"fmt"
    19  	"time"
    20  
    21  	. "github.com/insionng/yougam/libraries/pingcap/check"
    22  	"github.com/insionng/yougam/libraries/pingcap/tidb/util/codec"
    23  )
    24  
    25  var (
    26  	withTiKV     = flag.Bool("with-tikv", false, "run tests with TiKV cluster started. (not use the mock server)")
    27  	etcdAddrs    = flag.String("etcd-addrs", "127.0.0.1:2379", "etcd addrs")
    28  	pdLeaderPath = flag.String("pd-path", "/pd", "PD leader path on etcd")
    29  	clusterID    = flag.Int("cluster", 1, "cluster ID")
    30  )
    31  
    32  func newTestStore(c *C) *tikvStore {
    33  	if !flag.Parsed() {
    34  		flag.Parse()
    35  	}
    36  
    37  	if *withTiKV {
    38  		var d Driver
    39  		store, err := d.Open(fmt.Sprintf("tikv://%s%s?cluster=%d", *etcdAddrs, *pdLeaderPath, *clusterID))
    40  		c.Assert(err, IsNil)
    41  		return store.(*tikvStore)
    42  	}
    43  	return NewMockTikvStore().(*tikvStore)
    44  }
    45  
    46  type testTiclientSuite struct {
    47  	store *tikvStore
    48  	// prefix is prefix of each key in this test. It is used for table isolation,
    49  	// or it may pollute other data.
    50  	prefix string
    51  }
    52  
    53  var _ = Suite(&testTiclientSuite{})
    54  
    55  func (s *testTiclientSuite) SetUpSuite(c *C) {
    56  	s.store = newTestStore(c)
    57  	s.prefix = fmt.Sprintf("ticlient_%d", time.Now().Unix())
    58  }
    59  
    60  func (s *testTiclientSuite) TearDownSuite(c *C) {
    61  	// Clean all data, or it may pollute other data.
    62  	txn := s.beginTxn(c)
    63  	scanner, err := txn.Seek(encodeKey(s.prefix, ""))
    64  	c.Assert(err, IsNil)
    65  	c.Assert(scanner, NotNil)
    66  	for scanner.Valid() {
    67  		k := scanner.Key()
    68  		err = txn.Delete(k)
    69  		c.Assert(err, IsNil)
    70  		scanner.Next()
    71  	}
    72  	err = txn.Commit()
    73  	c.Assert(err, IsNil)
    74  	err = s.store.Close()
    75  	c.Assert(err, IsNil)
    76  }
    77  
    78  func (s *testTiclientSuite) beginTxn(c *C) *tikvTxn {
    79  	txn, err := s.store.Begin()
    80  	c.Assert(err, IsNil)
    81  	return txn.(*tikvTxn)
    82  }
    83  
    84  func (s *testTiclientSuite) TestSingleKey(c *C) {
    85  	txn := s.beginTxn(c)
    86  	err := txn.Set(encodeKey(s.prefix, "key"), []byte("value"))
    87  	c.Assert(err, IsNil)
    88  	err = txn.LockKeys(encodeKey(s.prefix, "key"))
    89  	c.Assert(err, IsNil)
    90  	err = txn.Commit()
    91  	c.Assert(err, IsNil)
    92  
    93  	txn = s.beginTxn(c)
    94  	val, err := txn.Get(encodeKey(s.prefix, "key"))
    95  	c.Assert(err, IsNil)
    96  	c.Assert(val, BytesEquals, []byte("value"))
    97  
    98  	txn = s.beginTxn(c)
    99  	err = txn.Delete(encodeKey(s.prefix, "key"))
   100  	c.Assert(err, IsNil)
   101  	err = txn.Commit()
   102  	c.Assert(err, IsNil)
   103  }
   104  
   105  func (s *testTiclientSuite) TestMultiKeys(c *C) {
   106  	const keyNum = 100
   107  
   108  	txn := s.beginTxn(c)
   109  	for i := 0; i < keyNum; i++ {
   110  		err := txn.Set(encodeKey(s.prefix, s08d("key", i)), valueBytes(i))
   111  		c.Assert(err, IsNil)
   112  	}
   113  	err := txn.Commit()
   114  	c.Assert(err, IsNil)
   115  
   116  	txn = s.beginTxn(c)
   117  	for i := 0; i < keyNum; i++ {
   118  		val, err1 := txn.Get(encodeKey(s.prefix, s08d("key", i)))
   119  		c.Assert(err1, IsNil)
   120  		c.Assert(val, BytesEquals, valueBytes(i))
   121  	}
   122  
   123  	txn = s.beginTxn(c)
   124  	for i := 0; i < keyNum; i++ {
   125  		err = txn.Delete(encodeKey(s.prefix, s08d("key", i)))
   126  		c.Assert(err, IsNil)
   127  	}
   128  	err = txn.Commit()
   129  	c.Assert(err, IsNil)
   130  }
   131  
   132  func (s *testTiclientSuite) TestCleanLock(c *C) {
   133  	const keyNum = 10
   134  
   135  	txn := s.beginTxn(c)
   136  	for i := 0; i < keyNum; i++ {
   137  		err := txn.Set(encodeKey(s.prefix, s08d("key", i)), valueBytes(i))
   138  		c.Assert(err, IsNil)
   139  	}
   140  	txn.DONOTCOMMIT = true
   141  	err := txn.Commit()
   142  	c.Assert(err, IsNil)
   143  
   144  	txn2 := s.beginTxn(c)
   145  	for i := 0; i < keyNum; i++ {
   146  		err2 := txn2.Set(encodeKey(s.prefix, s08d("key", i)), valueBytes(i+1))
   147  		c.Assert(err2, IsNil)
   148  	}
   149  	err2 := txn2.Commit()
   150  	c.Assert(err2, IsNil)
   151  }
   152  
   153  func (s *testTiclientSuite) TestNotExist(c *C) {
   154  	txn := s.beginTxn(c)
   155  	_, err := txn.Get(encodeKey(s.prefix, "noSuchKey"))
   156  	c.Assert(err, NotNil)
   157  }
   158  
   159  func encodeKey(prefix, s string) []byte {
   160  	return codec.EncodeBytes(nil, []byte(fmt.Sprintf("%s_%s", prefix, s)))
   161  }
   162  
   163  func valueBytes(n int) []byte {
   164  	return []byte(fmt.Sprintf("value%d", n))
   165  }
   166  
   167  // s08d is for returning format string "%s%08d" to keep string sorted.
   168  // e.g.: "0002" < "0011", otherwise "2" > "11"
   169  func s08d(prefix string, n int) string {
   170  	return fmt.Sprintf("%s%08d", prefix, n)
   171  }