Better OUTPUT tests

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)
 	}
 }

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":

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
 	}