github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/pingcap/tidb/store/tikv/client_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  	"net"
    18  	"testing"
    19  
    20  	"github.com/insionng/yougam/libraries/ngaut/log"
    21  	. "github.com/insionng/yougam/libraries/pingcap/check"
    22  	pb "github.com/insionng/yougam/libraries/pingcap/kvproto/pkg/kvrpcpb"
    23  	"github.com/insionng/yougam/libraries/pingcap/kvproto/pkg/msgpb"
    24  	"github.com/insionng/yougam/libraries/pingcap/kvproto/pkg/util"
    25  )
    26  
    27  func TestT(t *testing.T) {
    28  	TestingT(t)
    29  }
    30  
    31  var _ = Suite(&testClientSuite{})
    32  
    33  type testClientSuite struct {
    34  }
    35  
    36  // handleRequest receive Request then send empty Response back fill with same Type.
    37  func handleRequest(conn net.Conn, c *C) {
    38  	c.Assert(conn, NotNil)
    39  	defer conn.Close()
    40  	var msg msgpb.Message
    41  	msgID, err := util.ReadMessage(conn, &msg)
    42  	c.Assert(err, IsNil)
    43  	c.Assert(msgID, Greater, uint64(0))
    44  	c.Assert(msg.GetMsgType(), Equals, msgpb.MessageType_KvReq)
    45  
    46  	req := msg.GetKvReq()
    47  	c.Assert(req, NotNil)
    48  	var resp pb.Response
    49  	resp.Type = req.Type.Enum()
    50  	msg = msgpb.Message{
    51  		MsgType: msgpb.MessageType_KvResp.Enum(),
    52  		KvResp:  &resp,
    53  	}
    54  	err = util.WriteMessage(conn, msgID, &msg)
    55  	c.Assert(err, IsNil)
    56  }
    57  
    58  // One normally `Send`.
    59  func (s *testClientSuite) TestSendBySelf(c *C) {
    60  	l := startServer(":61234", c, handleRequest)
    61  	defer l.Close()
    62  	cli, err := NewRPCClient(":61234")
    63  	c.Assert(err, IsNil)
    64  	req := new(pb.Request)
    65  	req.Type = pb.MessageType_CmdGet.Enum()
    66  	getReq := new(pb.CmdGetRequest)
    67  	getReq.Key = []byte("a")
    68  	ver := uint64(0)
    69  	getReq.Version = &ver
    70  	req.CmdGetReq = getReq
    71  	resp, err := cli.SendKVReq(req)
    72  	c.Assert(err, IsNil)
    73  	c.Assert(req.GetType(), Equals, resp.GetType())
    74  }
    75  
    76  func closeRequest(conn net.Conn, c *C) {
    77  	c.Assert(conn, NotNil)
    78  	err := conn.Close()
    79  	c.Assert(err, IsNil)
    80  }
    81  
    82  // Server close connection directly if new connection is comming.
    83  func (s *testClientSuite) TestRetryClose(c *C) {
    84  	l := startServer(":61235", c, closeRequest)
    85  	defer l.Close()
    86  	cli, err := NewRPCClient(":61235")
    87  	c.Assert(err, IsNil)
    88  	req := new(pb.Request)
    89  	resp, err := cli.SendKVReq(req)
    90  	c.Assert(err, NotNil)
    91  	c.Assert(resp, IsNil)
    92  }
    93  
    94  func readThenCloseRequest(conn net.Conn, c *C) {
    95  	c.Assert(conn, NotNil)
    96  	defer conn.Close()
    97  	var msg msgpb.Message
    98  	msgID, err := util.ReadMessage(conn, &msg)
    99  	c.Assert(err, IsNil)
   100  	c.Assert(msg.GetKvReq(), NotNil)
   101  	c.Assert(msgID, Greater, uint64(0))
   102  }
   103  
   104  // Server read message then close, so `Send` will return retry error.
   105  func (s *testClientSuite) TestRetryReadThenClose(c *C) {
   106  	l := startServer(":61236", c, readThenCloseRequest)
   107  	defer l.Close()
   108  	cli, err := NewRPCClient(":61236")
   109  	c.Assert(err, IsNil)
   110  	req := new(pb.Request)
   111  	req.Type = pb.MessageType_CmdGet.Enum()
   112  	resp, err := cli.SendKVReq(req)
   113  	c.Assert(err, NotNil)
   114  	c.Assert(resp, IsNil)
   115  }
   116  
   117  func startServer(host string, c *C, handleFunc func(net.Conn, *C)) net.Listener {
   118  	l, err := net.Listen("tcp", host)
   119  	c.Assert(err, IsNil)
   120  	log.Debug("Start listenning on", host)
   121  	go func() {
   122  		for {
   123  			conn, err := l.Accept()
   124  			if err != nil {
   125  				return
   126  			}
   127  			go handleFunc(conn, c)
   128  		}
   129  	}()
   130  	return l
   131  }