github.com/hdt3213/godis@v1.2.9/database/set_test.go (about)

     1  package database
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/hdt3213/godis/lib/utils"
     6  	"github.com/hdt3213/godis/redis/protocol"
     7  	"github.com/hdt3213/godis/redis/protocol/asserts"
     8  	"math/rand"
     9  	"strconv"
    10  	"testing"
    11  )
    12  
    13  // basic add get and remove
    14  func TestSAdd(t *testing.T) {
    15  	testDB.Flush()
    16  	size := 100
    17  
    18  	// test sadd
    19  	key := utils.RandString(10)
    20  	for i := 0; i < size; i++ {
    21  		member := strconv.Itoa(i)
    22  		result := testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
    23  		asserts.AssertIntReply(t, result, 1)
    24  	}
    25  	// test scard
    26  	result := testDB.Exec(nil, utils.ToCmdLine("SCard", key))
    27  	asserts.AssertIntReply(t, result, size)
    28  
    29  	// test is member
    30  	for i := 0; i < size; i++ {
    31  		member := strconv.Itoa(i)
    32  		result = testDB.Exec(nil, utils.ToCmdLine("SIsMember", key, member))
    33  		asserts.AssertIntReply(t, result, 1)
    34  	}
    35  
    36  	// test members
    37  	result = testDB.Exec(nil, utils.ToCmdLine("SMembers", key))
    38  	multiBulk, ok := result.(*protocol.MultiBulkReply)
    39  	if !ok {
    40  		t.Error(fmt.Sprintf("expected bulk protocol, actually %s", result.ToBytes()))
    41  		return
    42  	}
    43  	if len(multiBulk.Args) != size {
    44  		t.Error(fmt.Sprintf("expected %d elements, actually %d", size, len(multiBulk.Args)))
    45  		return
    46  	}
    47  }
    48  
    49  func TestSRem(t *testing.T) {
    50  	testDB.Flush()
    51  	size := 100
    52  
    53  	// mock data
    54  	key := utils.RandString(10)
    55  	for i := 0; i < size; i++ {
    56  		member := strconv.Itoa(i)
    57  		testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
    58  	}
    59  	for i := 0; i < size; i++ {
    60  		member := strconv.Itoa(i)
    61  		testDB.Exec(nil, utils.ToCmdLine("srem", key, member))
    62  		result := testDB.Exec(nil, utils.ToCmdLine("SIsMember", key, member))
    63  		asserts.AssertIntReply(t, result, 0)
    64  	}
    65  }
    66  
    67  func TestSPop(t *testing.T) {
    68  	testDB.Flush()
    69  	size := 100
    70  
    71  	// mock data
    72  	key := utils.RandString(10)
    73  	for i := 0; i < size; i++ {
    74  		member := strconv.Itoa(i)
    75  		testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
    76  	}
    77  
    78  	result := testDB.Exec(nil, utils.ToCmdLine("spop", key))
    79  	asserts.AssertMultiBulkReplySize(t, result, 1)
    80  
    81  	currentSize := size - 1
    82  	for currentSize > 0 {
    83  		count := rand.Intn(currentSize) + 1
    84  		resultSpop := testDB.Exec(nil, utils.ToCmdLine("spop", key, strconv.FormatInt(int64(count), 10)))
    85  		multiBulk, ok := resultSpop.(*protocol.MultiBulkReply)
    86  		if !ok {
    87  			t.Error(fmt.Sprintf("expected bulk protocol, actually %s", resultSpop.ToBytes()))
    88  			return
    89  		}
    90  		removedSize := len(multiBulk.Args)
    91  		for _, arg := range multiBulk.Args {
    92  			resultSIsMember := testDB.Exec(nil, utils.ToCmdLine("SIsMember", key, string(arg)))
    93  			asserts.AssertIntReply(t, resultSIsMember, 0)
    94  		}
    95  		currentSize -= removedSize
    96  		resultSCard := testDB.Exec(nil, utils.ToCmdLine("SCard", key))
    97  		asserts.AssertIntReply(t, resultSCard, currentSize)
    98  	}
    99  }
   100  
   101  func TestSInter(t *testing.T) {
   102  	testDB.Flush()
   103  	size := 100
   104  	step := 10
   105  
   106  	keys := make([]string, 0)
   107  	start := 0
   108  	for i := 0; i < 4; i++ {
   109  		key := utils.RandString(10) + strconv.Itoa(i)
   110  		keys = append(keys, key)
   111  		for j := start; j < size+start; j++ {
   112  			member := strconv.Itoa(j)
   113  			testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
   114  		}
   115  		start += step
   116  	}
   117  	result := testDB.Exec(nil, utils.ToCmdLine2("sinter", keys...))
   118  	asserts.AssertMultiBulkReplySize(t, result, 70)
   119  
   120  	destKey := utils.RandString(10)
   121  	keysWithDest := []string{destKey}
   122  	keysWithDest = append(keysWithDest, keys...)
   123  	result = testDB.Exec(nil, utils.ToCmdLine2("SInterStore", keysWithDest...))
   124  	asserts.AssertIntReply(t, result, 70)
   125  
   126  	// test empty set
   127  	testDB.Flush()
   128  	key0 := utils.RandString(10)
   129  	testDB.Remove(key0)
   130  	key1 := utils.RandString(10)
   131  	testDB.Exec(nil, utils.ToCmdLine("sadd", key1, "a", "b"))
   132  	key2 := utils.RandString(10)
   133  	testDB.Exec(nil, utils.ToCmdLine("sadd", key1, "1", "2"))
   134  	result = testDB.Exec(nil, utils.ToCmdLine("sinter", key0, key1, key2))
   135  	asserts.AssertMultiBulkReplySize(t, result, 0)
   136  	result = testDB.Exec(nil, utils.ToCmdLine("sinter", key1, key2))
   137  	asserts.AssertMultiBulkReplySize(t, result, 0)
   138  	result = testDB.Exec(nil, utils.ToCmdLine("sinterstore", utils.RandString(10), key0, key1, key2))
   139  	asserts.AssertIntReply(t, result, 0)
   140  	result = testDB.Exec(nil, utils.ToCmdLine("sinterstore", utils.RandString(10), key1, key2))
   141  	asserts.AssertIntReply(t, result, 0)
   142  }
   143  
   144  func TestSUnion(t *testing.T) {
   145  	testDB.Flush()
   146  	size := 100
   147  	step := 10
   148  
   149  	keys := make([]string, 0)
   150  	start := 0
   151  	for i := 0; i < 4; i++ {
   152  		key := utils.RandString(10)
   153  		keys = append(keys, key)
   154  		for j := start; j < size+start; j++ {
   155  			member := strconv.Itoa(j)
   156  			testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
   157  		}
   158  		start += step
   159  	}
   160  	result := testDB.Exec(nil, utils.ToCmdLine2("sunion", keys...))
   161  	asserts.AssertMultiBulkReplySize(t, result, 130)
   162  
   163  	destKey := utils.RandString(10)
   164  	keysWithDest := []string{destKey}
   165  	keysWithDest = append(keysWithDest, keys...)
   166  	result = testDB.Exec(nil, utils.ToCmdLine2("SUnionStore", keysWithDest...))
   167  	asserts.AssertIntReply(t, result, 130)
   168  }
   169  
   170  func TestSDiff(t *testing.T) {
   171  	testDB.Flush()
   172  	size := 100
   173  	step := 20
   174  
   175  	keys := make([]string, 0)
   176  	start := 0
   177  	for i := 0; i < 3; i++ {
   178  		key := utils.RandString(10)
   179  		keys = append(keys, key)
   180  		for j := start; j < size+start; j++ {
   181  			member := strconv.Itoa(j)
   182  			testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
   183  		}
   184  		start += step
   185  	}
   186  	result := testDB.Exec(nil, utils.ToCmdLine2("SDiff", keys...))
   187  	asserts.AssertMultiBulkReplySize(t, result, step)
   188  
   189  	destKey := utils.RandString(10)
   190  	keysWithDest := []string{destKey}
   191  	keysWithDest = append(keysWithDest, keys...)
   192  	result = testDB.Exec(nil, utils.ToCmdLine2("SDiffStore", keysWithDest...))
   193  	asserts.AssertIntReply(t, result, step)
   194  
   195  	// test empty set
   196  	testDB.Flush()
   197  	key0 := utils.RandString(10)
   198  	testDB.Remove(key0)
   199  	key1 := utils.RandString(10)
   200  	testDB.Exec(nil, utils.ToCmdLine("sadd", key1, "a", "b"))
   201  	key2 := utils.RandString(10)
   202  	testDB.Exec(nil, utils.ToCmdLine("sadd", key2, "a", "b"))
   203  	result = testDB.Exec(nil, utils.ToCmdLine("sdiff", key0, key1, key2))
   204  	asserts.AssertMultiBulkReplySize(t, result, 0)
   205  	result = testDB.Exec(nil, utils.ToCmdLine("sdiff", key1, key2))
   206  	asserts.AssertMultiBulkReplySize(t, result, 0)
   207  	result = testDB.Exec(nil, utils.ToCmdLine("SDiffStore", utils.RandString(10), key0, key1, key2))
   208  	asserts.AssertIntReply(t, result, 0)
   209  	result = testDB.Exec(nil, utils.ToCmdLine("SDiffStore", utils.RandString(10), key1, key2))
   210  	asserts.AssertIntReply(t, result, 0)
   211  }
   212  
   213  func TestSRandMember(t *testing.T) {
   214  	testDB.Flush()
   215  	key := utils.RandString(10)
   216  	for j := 0; j < 100; j++ {
   217  		member := strconv.Itoa(j)
   218  		testDB.Exec(nil, utils.ToCmdLine("sadd", key, member))
   219  	}
   220  	result := testDB.Exec(nil, utils.ToCmdLine("SRandMember", key))
   221  	br, ok := result.(*protocol.BulkReply)
   222  	if !ok && len(br.Arg) > 0 {
   223  		t.Error(fmt.Sprintf("expected bulk protocol, actually %s", result.ToBytes()))
   224  		return
   225  	}
   226  
   227  	result = testDB.Exec(nil, utils.ToCmdLine("SRandMember", key, "10"))
   228  	asserts.AssertMultiBulkReplySize(t, result, 10)
   229  	multiBulk, ok := result.(*protocol.MultiBulkReply)
   230  	if !ok {
   231  		t.Error(fmt.Sprintf("expected bulk protocol, actually %s", result.ToBytes()))
   232  		return
   233  	}
   234  	m := make(map[string]struct{})
   235  	for _, arg := range multiBulk.Args {
   236  		m[string(arg)] = struct{}{}
   237  	}
   238  	if len(m) != 10 {
   239  		t.Error(fmt.Sprintf("expected 10 members, actually %d", len(m)))
   240  		return
   241  	}
   242  
   243  	result = testDB.Exec(nil, utils.ToCmdLine("SRandMember", key, "110"))
   244  	asserts.AssertMultiBulkReplySize(t, result, 100)
   245  
   246  	result = testDB.Exec(nil, utils.ToCmdLine("SRandMember", key, "-10"))
   247  	asserts.AssertMultiBulkReplySize(t, result, 10)
   248  
   249  	result = testDB.Exec(nil, utils.ToCmdLine("SRandMember", key, "-110"))
   250  	asserts.AssertMultiBulkReplySize(t, result, 110)
   251  }