Merge pull request #158 from alicebob/seterror

introduce SetError()

Author: alicebob

cmd_connection_test.go

@@ -170,3 +170,24 @@ func TestQuit(t *testing.T) {
 	assert(t, err != nil, "QUIT closed the client")
 	equals(t, "", v)
 }
+
+func TestSetError(t *testing.T) {
+	s, err := Run()
+	ok(t, err)
+	defer s.Close()
+	c, err := redis.Dial("tcp", s.Addr())
+	ok(t, err)
+
+	r, err := redis.String(c.Do("PING"))
+	ok(t, err)
+	equals(t, "PONG", r)
+
+	s.SetError("LOADING Redis is loading the dataset in memory")
+	_, err = c.Do("PING", "hi")
+	mustFail(t, err, "LOADING Redis is loading the dataset in memory")
+
+	s.SetError("")
+	r, err = redis.String(c.Do("PING"))
+	ok(t, err)
+	equals(t, "PONG", r)
+}

miniredis.go

@@ -354,6 +354,21 @@ func (m *Miniredis) SetTime(t time.Time) {
 	m.now = t
 }
 
+// make every command return this message. For example:
+//   LOADING Redis is loading the dataset in memory
+//   MASTERDOWN Link with MASTER is down and replica-serve-stale-data is set to 'no'.
+// Clear it with an empty string. Don't add newlines.
+func (m *Miniredis) SetError(msg string) {
+	cb := server.Hook(nil)
+	if msg != "" {
+		cb = func(c *server.Peer, cmd string, args ...string) bool {
+			c.WriteError(msg)
+			return true
+		}
+	}
+	m.srv.SetPreHook(cb)
+}
+
 // handleAuth returns false if connection has no access. It sends the reply.
 func (m *Miniredis) handleAuth(c *server.Peer) bool {
 	if getCtx(c).nested {

server/server.go

@@ -25,10 +25,14 @@ type Cmd func(c *Peer, cmd string, args []string)
 
 type DisconnectHandler func(c *Peer)
 
+// Hook is can be added to run before every cmd. Return true if the command is done.
+type Hook func(*Peer, string, ...string) bool
+
 // Server is a simple redis server
 type Server struct {
 	l         net.Listener
 	cmds      map[string]Cmd
+	preHook   Hook
 	peers     map[net.Conn]struct{}
 	mu        sync.Mutex
 	wg        sync.WaitGroup
@@ -63,6 +67,13 @@ func NewServer(addr string) (*Server, error) {
 	return &s, nil
 }
 
+// (un)set a hook which is ran before every call. It returns true if the command is done.
+func (s *Server) SetPreHook(h Hook) {
+	s.mu.Lock()
+	s.preHook = h
+	s.mu.Unlock()
+}
+
 func (s *Server) serve(l net.Listener) {
 	for {
 		conn, err := l.Accept()
@@ -160,6 +171,15 @@ func (s *Server) servePeer(c net.Conn) {
 func (s *Server) Dispatch(c *Peer, args []string) {
 	cmd, args := args[0], args[1:]
 	cmdUp := strings.ToUpper(cmd)
+	s.mu.Lock()
+	h := s.preHook
+	s.mu.Unlock()
+	if h != nil {
+		if h(c, cmdUp, args...) {
+			return
+		}
+	}
+
 	s.mu.Lock()
 	cb, ok := s.cmds[cmdUp]
 	s.mu.Unlock()