Merge 5ac42cd into 2be4728

command/init.go

@@ -28,11 +28,12 @@ func init() {
 		"getset":   GetSet,
 		"getrange": GetRange,
 		// "msetnx":   MSetNx,
-		"setnx":    SetNx,
-		"setex":    SetEx,
-		"psetex":   PSetEx,
-		"setrange": SetRange,
-		// "setbit":      SetBit,
+		"setnx":       SetNx,
+		"setex":       SetEx,
+		"psetex":      PSetEx,
+		"setrange":    SetRange,
+		"setbit":      SetBit,
+		"getbit":      GetBit,
 		"incr":        Incr,
 		"incrby":      IncrBy,
 		"decr":        Decr,
@@ -129,6 +130,8 @@ func init() {
 		"incrby":      Desc{Proc: AutoCommit(IncrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"decrby":      Desc{Proc: AutoCommit(DecrBy), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
 		"incrbyfloat": Desc{Proc: AutoCommit(IncrByFloat), Cons: Constraint{3, flags("wmF"), 1, 1, 1}},
+		"setbit":      Desc{Proc: AutoCommit(SetBit), Cons: Constraint{4, flags("wm"), 1, 1, 1}},
+		"getbit":      Desc{Proc: AutoCommit(GetBit), Cons: Constraint{3, flags("r"), 1, 1, 1}},
 
 		// keys
 		"type":      Desc{Proc: AutoCommit(Type), Cons: Constraint{2, flags("rF"), 1, 1, 1}},

command/strings.go

@@ -469,3 +469,68 @@ func DecrBy(ctx *Context, txn *db.Transaction) (OnCommit, error) {
 	}
 	return Integer(ctx.Out, int64(delta)), nil
 }
+
+// SetBit Sets or clears the bit at offset in the string value stored at key.
+func SetBit(ctx *Context, txn *db.Transaction) (OnCommit, error) {
+	key := []byte(ctx.Args[0])
+	offset, err := strconv.Atoi(string(ctx.Args[1]))
+	if err != nil {
+		return nil, ErrBitOffset
+	}
+	if offset < 0 {
+		return nil, ErrBitOffset
+	}
+
+	on, err := strconv.Atoi(string(ctx.Args[2]))
+	if err != nil {
+		return nil, ErrBitInteger
+	}
+
+	// Bits can only be set or cleared...
+	if (on & ^1) != 0 {
+		return nil, ErrBitInteger
+	}
+
+	str, err := txn.String(key)
+	if err != nil {
+		if err == db.ErrTypeMismatch {
+			return nil, ErrTypeMismatch
+		}
+		return nil, errors.New("ERR " + err.Error())
+	}
+	val, err := str.SetBit(offset, on)
+	if err != nil {
+		return nil, errors.New("ERR " + err.Error())
+	}
+	if val != 0 {
+		return Integer(ctx.Out, 1), nil
+	}
+	return Integer(ctx.Out, 0), nil
+}
+
+// GetBit get the bit at offset in the string value stored at key.
+func GetBit(ctx *Context, txn *db.Transaction) (OnCommit, error) {
+	key := []byte(ctx.Args[0])
+	offset, err := strconv.Atoi(string(ctx.Args[1]))
+	if err != nil || offset < 0 {
+		return nil, ErrBitOffset
+	}
+
+	str, err := txn.String(key)
+	if err != nil {
+		if err == db.ErrTypeMismatch {
+			return nil, ErrTypeMismatch
+		}
+		return nil, errors.New("ERR " + err.Error())
+	}
+
+	val, err := str.GetBit(offset)
+	if err != nil {
+		return nil, errors.New("ERR " + err.Error())
+	}
+
+	if val != 0 {
+		return Integer(ctx.Out, 1), nil
+	}
+	return Integer(ctx.Out, 0), nil
+}

command/strings_test.go

@@ -497,11 +497,97 @@ func TestStringAppend(t *testing.T) {
 	args[0] = "Append"
 	args[1] = "Ap"
 
-	ctx := ContextTest("append", args...)
-	Call(ctx)
-	assert.Contains(t, ctxString(ctx.Out), strconv.Itoa(len(args[1])))
+	out := CallTest("append", args...)
+	assert.Contains(t, out.String(), strconv.Itoa(len(args[1])))
 
-	ctx = ContextTest("append", args...)
-	Call(ctx)
-	assert.Contains(t, ctxString(ctx.Out), strconv.Itoa(len(args[1])*2))
+	out = CallTest("append", args...)
+	assert.Contains(t, out.String(), strconv.Itoa(len(args[1])*2))
+}
+
+func TestStringSetBit(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "1",
+			args: []string{"setbit", "1", "0", "0"},
+			want: "-ERR wrong number of arguments for 'setbit' command",
+		},
+		{
+			name: "2",
+			args: []string{"setbit", "x", "0"},
+			want: ErrBitOffset.Error(),
+		},
+		{
+			name: "3",
+			args: []string{"setbit", "1", "x"},
+			want: ErrBitInteger.Error(),
+		},
+		{
+			name: "4",
+			args: []string{"setbit", "1", "2"},
+			want: ErrBitInteger.Error(),
+		},
+		{
+			name: "5",
+			args: []string{"setbit", "1", "1"},
+			want: ":0",
+		},
+		{
+			name: "6",
+			args: []string{"setbit", "1", "0"},
+			want: ":1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			out := CallTest("setbit", tt.args...)
+			assert.Contains(t, out.String(), tt.want)
+		})
+	}
+}
+
+func TestStringGetBit(t *testing.T) {
+	CallTest("setbit", "getbit", "5", "1")
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "1",
+			args: []string{"getbit", "1", "0", "0"},
+			want: "-ERR wrong number of arguments for 'getbit' command",
+		},
+		{
+			name: "2",
+			args: []string{"getbit", "x"},
+			want: ErrBitOffset.Error(),
+		},
+		{
+			name: "3",
+			args: []string{"getbit", "1"},
+			want: ":0",
+		},
+		{
+			name: "5",
+			args: []string{"getbit", "5"},
+			want: ":1",
+		},
+		{
+			name: "6",
+			args: []string{"getbit", "10"},
+			want: ":0",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			out := CallTest("getbit", tt.args...)
+			assert.Contains(t, out.String(), tt.want)
+		})
+	}
 }

db/db_test.go

@@ -1,11 +1,13 @@
 package db
 
 import (
+	"context"
 	"os"
 	"testing"
 
 	"github.com/meitu/titan/conf"
 	"github.com/pingcap/tidb/store/mockstore"
+	"github.com/stretchr/testify/assert"
 )
 
 var mockDB *DB
@@ -25,3 +27,11 @@ func TestMain(m *testing.M) {
 
 	os.Exit(m.Run())
 }
+
+func MockTest(t *testing.T, callFunc func(txn *Transaction)) {
+	txn, err := mockDB.Begin()
+	assert.NoError(t, err)
+	callFunc(txn)
+	err = txn.Commit(context.TODO())
+	assert.NoError(t, err)
+}