github.com/hdt3213/godis@v1.2.9/cluster/utils_test.go (about)

     1  package cluster
     2  
     3  import (
     4  	"github.com/hdt3213/godis/config"
     5  	"github.com/hdt3213/godis/interface/redis"
     6  	"github.com/hdt3213/godis/redis/protocol"
     7  	"math/rand"
     8  	"strings"
     9  )
    10  
    11  var testNodeA, testNodeB *Cluster
    12  var simulateATimout, simulateBTimout *bool
    13  
    14  type mockPicker struct {
    15  	nodes []string
    16  }
    17  
    18  func (picker *mockPicker) AddNode(keys ...string) {
    19  	picker.nodes = append(picker.nodes, keys...)
    20  }
    21  
    22  func (picker *mockPicker) PickNode(key string) string {
    23  	for _, n := range picker.nodes {
    24  		if strings.Contains(key, n) {
    25  			return n
    26  		}
    27  	}
    28  	return picker.nodes[0]
    29  }
    30  
    31  func makeMockRelay(peer *Cluster) (*bool, func(cluster *Cluster, node string, c redis.Connection, cmdLine CmdLine) redis.Reply) {
    32  	simulateTimeout0 := false
    33  	simulateTimeout := &simulateTimeout0
    34  	return simulateTimeout, func(cluster *Cluster, node string, c redis.Connection, cmdLine CmdLine) redis.Reply {
    35  		if len(cmdLine) == 0 {
    36  			return protocol.MakeErrReply("ERR command required")
    37  		}
    38  		if node == cluster.self {
    39  			// to self db
    40  			cmdName := strings.ToLower(string(cmdLine[0]))
    41  			if cmdName == "prepare" {
    42  				return execPrepare(cluster, c, cmdLine)
    43  			} else if cmdName == "commit" {
    44  				return execCommit(cluster, c, cmdLine)
    45  			} else if cmdName == "rollback" {
    46  				return execRollback(cluster, c, cmdLine)
    47  			}
    48  			return cluster.db.Exec(c, cmdLine)
    49  		}
    50  		if *simulateTimeout {
    51  			return protocol.MakeErrReply("ERR timeout")
    52  		}
    53  		cmdName := strings.ToLower(string(cmdLine[0]))
    54  		if cmdName == "prepare" {
    55  			return execPrepare(peer, c, cmdLine)
    56  		} else if cmdName == "commit" {
    57  			return execCommit(peer, c, cmdLine)
    58  		} else if cmdName == "rollback" {
    59  			return execRollback(peer, c, cmdLine)
    60  		}
    61  		return peer.db.Exec(c, cmdLine)
    62  	}
    63  }
    64  
    65  func init() {
    66  	if config.Properties == nil {
    67  		config.Properties = &config.ServerProperties{}
    68  	}
    69  	addrA := "127.0.0.1:6399"
    70  	addrB := "127.0.0.1:7379"
    71  	config.Properties.Self = addrA
    72  	config.Properties.Peers = []string{addrB}
    73  	testNodeA = MakeCluster()
    74  	config.Properties.Self = addrB
    75  	config.Properties.Peers = []string{addrA}
    76  	testNodeB = MakeCluster()
    77  
    78  	simulateBTimout, testNodeA.relayImpl = makeMockRelay(testNodeB)
    79  	testNodeA.peerPicker = &mockPicker{}
    80  	testNodeA.peerPicker.AddNode(addrA, addrB)
    81  	simulateATimout, testNodeB.relayImpl = makeMockRelay(testNodeA)
    82  	testNodeB.peerPicker = &mockPicker{}
    83  	testNodeB.peerPicker.AddNode(addrB, addrA)
    84  }
    85  
    86  func MakeTestCluster(peers []string) *Cluster {
    87  	if config.Properties == nil {
    88  		config.Properties = &config.ServerProperties{}
    89  	}
    90  	config.Properties.Self = "127.0.0.1:6399"
    91  	config.Properties.Peers = peers
    92  	return MakeCluster()
    93  }
    94  
    95  func toArgs(cmd ...string) [][]byte {
    96  	args := make([][]byte, len(cmd))
    97  	for i, s := range cmd {
    98  		args[i] = []byte(s)
    99  	}
   100  	return args
   101  }
   102  
   103  var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
   104  
   105  func RandString(n int) string {
   106  	b := make([]rune, n)
   107  	for i := range b {
   108  		b[i] = letters[rand.Intn(len(letters))]
   109  	}
   110  	return string(b)
   111  }