package valkey import ( "context" "crypto/sha1" "encoding/hex" "math/rand" "reflect" "strconv" "testing" "time" "github.com/valkey-io/valkey-go/internal/cmds" ) func TestNewLuaScriptOnePass(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) sum := sha1.Sum([]byte(body)) sha := hex.EncodeToString(sum[:]) k := []string{"1", "2"} a := []string{"3", "4"} c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { if reflect.DeepEqual(cmd.Commands(), []string{"EVALSHA", sha, "2", "1", "2", "3", "4"}) { return newResult(ValkeyMessage{typ: '+', string: "OK"}, nil) } return newResult(ValkeyMessage{typ: '+', string: "unexpected"}, nil) }, } script := NewLuaScript(body) if v, err := script.Exec(context.Background(), c, k, a).ToString(); err != nil || v != "OK" { t.Fatalf("ret mistmatch") } } func TestNewLuaScript(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) sum := sha1.Sum([]byte(body)) sha := hex.EncodeToString(sum[:]) k := []string{"1", "2"} a := []string{"3", "4"} eval := false c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { if reflect.DeepEqual(cmd.Commands(), []string{"EVALSHA", sha, "2", "1", "2", "3", "4"}) { eval = true return newResult(ValkeyMessage{typ: '-', string: "NOSCRIPT"}, nil) } if eval && reflect.DeepEqual(cmd.Commands(), []string{"EVAL", body, "2", "1", "2", "3", "4"}) { return newResult(ValkeyMessage{typ: '_'}, nil) } return newResult(ValkeyMessage{typ: '+', string: "unexpected"}, nil) }, } script := NewLuaScript(body) if err, ok := IsValkeyErr(script.Exec(context.Background(), c, k, a).Error()); ok && !err.IsNil() { t.Fatalf("ret mistmatch") } } func TestNewLuaScriptReadOnly(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) sum := sha1.Sum([]byte(body)) sha := hex.EncodeToString(sum[:]) k := []string{"1", "2"} a := []string{"3", "4"} eval := false c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { if reflect.DeepEqual(cmd.Commands(), []string{"EVALSHA_RO", sha, "2", "1", "2", "3", "4"}) { eval = true return newResult(ValkeyMessage{typ: '-', string: "NOSCRIPT"}, nil) } if eval && reflect.DeepEqual(cmd.Commands(), []string{"EVAL_RO", body, "2", "1", "2", "3", "4"}) { return newResult(ValkeyMessage{typ: '_'}, nil) } return newResult(ValkeyMessage{typ: '+', string: "unexpected"}, nil) }, } script := NewLuaScriptReadOnly(body) if err, ok := IsValkeyErr(script.Exec(context.Background(), c, k, a).Error()); ok && !err.IsNil() { t.Fatalf("ret mistmatch") } } func TestNewLuaScriptExecMultiError(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) k := []string{"1", "2"} a := []string{"3", "4"} c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { return newResult(ValkeyMessage{typ: '-', string: "ANY ERR"}, nil) }, } script := NewLuaScript(body) if script.ExecMulti(context.Background(), c, LuaExec{Keys: k, Args: a})[0].Error().Error() != "ANY ERR" { t.Fatalf("ret mistmatch") } } func TestNewLuaScriptExecMulti(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) sum := sha1.Sum([]byte(body)) sha := hex.EncodeToString(sum[:]) k := []string{"1", "2"} a := []string{"3", "4"} c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { return newResult(ValkeyMessage{typ: '+', string: "OK"}, nil) }, DoMultiFn: func(ctx context.Context, multi ...Completed) (resp []ValkeyResult) { for _, cmd := range multi { if reflect.DeepEqual(cmd.Commands(), []string{"EVALSHA", sha, "2", "1", "2", "3", "4"}) { resp = append(resp, newResult(ValkeyMessage{typ: '+', string: "OK"}, nil)) } } return resp }, } script := NewLuaScript(body) if v, err := script.ExecMulti(context.Background(), c, LuaExec{Keys: k, Args: a})[0].ToString(); err != nil || v != "OK" { t.Fatalf("ret mistmatch") } } func TestNewLuaScriptExecMultiRo(t *testing.T) { defer ShouldNotLeaked(SetupLeakDetection()) body := strconv.Itoa(rand.Int()) sum := sha1.Sum([]byte(body)) sha := hex.EncodeToString(sum[:]) k := []string{"1", "2"} a := []string{"3", "4"} c := &client{ BFn: func() Builder { return cmds.NewBuilder(cmds.NoSlot) }, DoFn: func(ctx context.Context, cmd Completed) (resp ValkeyResult) { return newResult(ValkeyMessage{typ: '+', string: "OK"}, nil) }, DoMultiFn: func(ctx context.Context, multi ...Completed) (resp []ValkeyResult) { for _, cmd := range multi { if reflect.DeepEqual(cmd.Commands(), []string{"EVALSHA_RO", sha, "2", "1", "2", "3", "4"}) { resp = append(resp, newResult(ValkeyMessage{typ: '+', string: "OK"}, nil)) } } return resp }, } script := NewLuaScriptReadOnly(body) if v, err := script.ExecMulti(context.Background(), c, LuaExec{Keys: k, Args: a})[0].ToString(); err != nil || v != "OK" { t.Fatalf("ret mistmatch") } } type client struct { BFn func() Builder DoFn func(ctx context.Context, cmd Completed) (resp ValkeyResult) DoMultiFn func(ctx context.Context, cmd ...Completed) (resp []ValkeyResult) DoCacheFn func(ctx context.Context, cmd Cacheable, ttl time.Duration) (resp ValkeyResult) DoMultiCacheFn func(ctx context.Context, cmd ...CacheableTTL) (resp []ValkeyResult) DedicatedFn func(fn func(DedicatedClient) error) (err error) DedicateFn func() (DedicatedClient, func()) CloseFn func() } func (c *client) Receive(ctx context.Context, subscribe Completed, fn func(msg PubSubMessage)) error { return nil } func (c *client) B() Builder { if c.BFn != nil { return c.BFn() } return Builder{} } func (c *client) Do(ctx context.Context, cmd Completed) (resp ValkeyResult) { if c.DoFn != nil { return c.DoFn(ctx, cmd) } return ValkeyResult{} } func (c *client) DoMulti(ctx context.Context, cmd ...Completed) (resp []ValkeyResult) { if c.DoMultiFn != nil { return c.DoMultiFn(ctx, cmd...) } return nil } func (c *client) DoStream(ctx context.Context, cmd Completed) (resp ValkeyResultStream) { return ValkeyResultStream{} } func (c *client) DoMultiStream(ctx context.Context, cmd ...Completed) (resp MultiValkeyResultStream) { return MultiValkeyResultStream{} } func (c *client) DoMultiCache(ctx context.Context, cmd ...CacheableTTL) (resp []ValkeyResult) { if c.DoMultiCacheFn != nil { return c.DoMultiCacheFn(ctx, cmd...) } return nil } func (c *client) DoCache(ctx context.Context, cmd Cacheable, ttl time.Duration) (resp ValkeyResult) { if c.DoCacheFn != nil { return c.DoCacheFn(ctx, cmd, ttl) } return ValkeyResult{} } func (c *client) Dedicated(fn func(DedicatedClient) error) (err error) { if c.DedicatedFn != nil { return c.DedicatedFn(fn) } return nil } func (c *client) Dedicate() (DedicatedClient, func()) { if c.DedicateFn != nil { return c.DedicateFn() } return nil, nil } func (c *client) Nodes() map[string]Client { return map[string]Client{"addr": c} } func (c *client) Close() { if c.CloseFn != nil { c.CloseFn() } } func ExampleLua_exec() { client, err := NewClient(ClientOption{InitAddress: []string{"127.0.0.1:6379"}}) if err != nil { panic(err) } defer client.Close() ctx := context.Background() script := NewLuaScript("return {KEYS[1],KEYS[2],ARGV[1],ARGV[2]}") script.Exec(ctx, client, []string{"k1", "k2"}, []string{"a1", "a2"}).ToArray() }