diff --git a/internal/server/output.go b/internal/server/output.go index 65824ada..51476a1e 100644 --- a/internal/server/output.go +++ b/internal/server/output.go @@ -7,35 +7,31 @@ import ( "github.com/tidwall/resp" ) -func (s *Server) cmdOutput(msg *Message) (res resp.Value, err error) { +// OUTPUT [resp|json] +func (s *Server) cmdOUTPUT(msg *Message) (resp.Value, error) { start := time.Now() - vs := msg.Args[1:] - var arg string - var ok bool - if len(vs) != 0 { - if _, arg, ok = tokenval(vs); !ok || arg == "" { - return NOMessage, errInvalidNumberOfArguments + args := msg.Args + switch len(args) { + case 1: + if msg.OutputType == JSON { + return resp.StringValue(`{"ok":true,"output":"json","elapsed":` + + time.Since(start).String() + `}`), nil } + return resp.StringValue("resp"), nil + case 2: // Setting the original message output type will be picked up by the // server prior to the next command being executed. - switch strings.ToLower(arg) { + switch strings.ToLower(args[1]) { default: - return NOMessage, errInvalidArgument(arg) + return retrerr(errInvalidArgument(args[1])) case "json": msg.OutputType = JSON case "resp": msg.OutputType = RESP } return OKMessage(msg, start), nil - } - // return the output - switch msg.OutputType { default: - return NOMessage, nil - case JSON: - return resp.StringValue(`{"ok":true,"output":"json","elapsed":` + time.Since(start).String() + `}`), nil - case RESP: - return resp.StringValue("resp"), nil + return retrerr(errInvalidNumberOfArguments) } } diff --git a/internal/server/server.go b/internal/server/server.go index 46865806..bf792724 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1215,7 +1215,7 @@ func (s *Server) command(msg *Message, client *Client) ( case "keys": res, err = s.cmdKEYS(msg) case "output": - res, err = s.cmdOutput(msg) + res, err = s.cmdOUTPUT(msg) case "aof": res, err = s.cmdAOF(msg) case "aofmd5": diff --git a/tests/client_test.go b/tests/client_test.go index d7d754fa..2c8d3b7d 100644 --- a/tests/client_test.go +++ b/tests/client_test.go @@ -18,8 +18,14 @@ func subTestClient(g *testGroup) { func client_OUTPUT_test(mc *mockServer) error { if err := mc.DoBatch( // tests removal of "elapsed" member. + Do("OUTPUT", "json", "yaml").Err(`wrong number of arguments for 'output' command`), Do("OUTPUT", "json").Str(`{"ok":true}`), + Do("OUTPUT").JSON().Str(`{"ok":true,"output":"json"}`), + Do("OUTPUT").Str(`resp`), // this is due to the internal Do test Do("OUTPUT", "resp").OK(), + Do("OUTPUT", "yaml").Err(`invalid argument 'yaml'`), + Do("OUTPUT").Str(`resp`), + Do("OUTPUT").JSON().Str(`{"ok":true,"output":"json"}`), ); err != nil { return err }