Add protocol to encode command

- Mutate response in encode when protocol is `gql`

README.md

@@ -154,6 +154,8 @@ encode command:
     	Output file (default "stdout")
   -to string
     	Output encoding [csv, gob, json] (default "json")
+  -protocol string
+    	Protocol of the results being encoded [http, gql] (default "http")
 
 plot command:
   -output string
@@ -700,8 +702,9 @@ Arguments:
           the supported encodings (gob | json | csv) [default: stdin]
 
 Options:
-  --to      Output encoding (gob | json | csv) [default: json]
-  --output  Output file [default: stdout]
+  --to        Output encoding (gob | json | csv) [default: json]
+  --output    Output file [default: stdout]
+  --protocol  Protocol of the results being encoded (http | gql) [default: gql]
 
 Examples:
   echo "GET http://:80" | vegeta attack -rate=1/s > results.gob

encode.go

@@ -44,8 +44,9 @@ Arguments:
           the supported encodings (gob | json | csv) [default: stdin]
 
 Options:
-  --to      Output encoding (gob | json | csv) [default: json]
-  --output  Output file [default: stdout]
+  --to        Output encoding (gob | json | csv) [default: json]
+  --output    Output file [default: stdout]
+  --protocol  Protocol of the results being encoded (http | gql) [default: http]
 
 Examples:
   echo "GET http://:80" | vegeta attack -rate=1/s > results.gob
@@ -57,6 +58,7 @@ func encodeCmd() command {
 	fs := flag.NewFlagSet("vegeta encode", flag.ExitOnError)
 	to := fs.String("to", encodingJSON, "Output encoding "+encs)
 	output := fs.String("output", "stdout", "Output file")
+	protocol := fs.String("protocol", "http", "Protocol of the results being encoded [http, gql]")
 
 	fs.Usage = func() {
 		fmt.Fprintf(os.Stderr, "%s\n", encodeUsage)
@@ -68,11 +70,11 @@ func encodeCmd() command {
 		if len(files) == 0 {
 			files = append(files, "stdin")
 		}
-		return encode(files, *to, *output)
+		return encode(files, *to, *output, *protocol)
 	}}
 }
 
-func encode(files []string, to, output string) error {
+func encode(files []string, to, output string, protocol string) error {
 	dec, mc, err := decoder(files)
 	defer mc.Close()
 	if err != nil {
@@ -113,7 +115,15 @@ func encode(files []string, to, output string) error {
 				break
 			}
 			return err
-		} else if err = enc.Encode(&r); err != nil {
+		}
+
+		if protocol == "gql" {
+			if err = vegeta.AsGraphQL(&r); err != nil {
+				return err
+			}
+		}
+
+		if err = enc.Encode(&r); err != nil {
 			return err
 		}
 	}

lib/results.go

@@ -6,6 +6,8 @@ import (
 	"encoding/base64"
 	"encoding/csv"
 	"encoding/gob"
+	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 	"net/textproto"
@@ -155,8 +157,8 @@ func (dec Decoder) Decode(r *Result) error { return dec(r) }
 type Encoder func(*Result) error
 
 // NewEncoder returns a new Result encoder closure for the given io.Writer
-func NewEncoder(r io.Writer) Encoder {
-	enc := gob.NewEncoder(r)
+func NewEncoder(w io.Writer) Encoder {
+	enc := gob.NewEncoder(w)
 	return func(r *Result) error { return enc.Encode(r) }
 }
 
@@ -205,8 +207,8 @@ func headerBytes(h http.Header) []byte {
 }
 
 // NewCSVDecoder returns a Decoder that decodes CSV encoded Results.
-func NewCSVDecoder(r io.Reader) Decoder {
-	dec := csv.NewReader(r)
+func NewCSVDecoder(rd io.Reader) Decoder {
+	dec := csv.NewReader(rd)
 	dec.FieldsPerRecord = 12
 	dec.TrimLeadingSpace = true
 
@@ -276,27 +278,62 @@ type jsonResult Result
 // NewJSONEncoder returns an Encoder that dumps the given *Results as a JSON
 // object.
 func NewJSONEncoder(w io.Writer) Encoder {
-	var jw jwriter.Writer
+	var enc jwriter.Writer
 	return func(r *Result) error {
-		(*jsonResult)(r).MarshalEasyJSON(&jw)
-		if jw.Error != nil {
-			return jw.Error
+		(*jsonResult)(r).MarshalEasyJSON(&enc)
+		if enc.Error != nil {
+			return enc.Error
 		}
-		jw.RawByte('\n')
-		_, err := jw.DumpTo(w)
+		enc.RawByte('\n')
+		_, err := enc.DumpTo(w)
 		return err
 	}
 }
 
 // NewJSONDecoder returns a Decoder that decodes JSON encoded Results.
-func NewJSONDecoder(r io.Reader) Decoder {
-	rd := bufio.NewReader(r)
+func NewJSONDecoder(rd io.Reader) Decoder {
+	dec := bufio.NewReader(rd)
 	return func(r *Result) (err error) {
 		var jl jlexer.Lexer
-		if jl.Data, err = rd.ReadBytes('\n'); err != nil {
+		if jl.Data, err = dec.ReadBytes('\n'); err != nil {
 			return err
 		}
 		(*jsonResult)(r).UnmarshalEasyJSON(&jl)
 		return jl.Error()
 	}
 }
+
+type GQLResponse struct {
+	Data   interface{}
+	Errors []GQLError
+}
+
+type GQLError struct {
+	Extensions interface{}
+	Message    string
+}
+
+// AsGraphQL re-interprets the given Result with GraphQL semantics, mutating it accordingly
+func AsGraphQL(r *Result) error {
+	if r.Code < 200 || r.Code >= 400 {
+		return nil
+	}
+
+	var res GQLResponse
+	err := json.Unmarshal(r.Body, &res)
+	if err != nil {
+		return err
+	}
+
+	if res.Errors != nil && len(res.Errors) > 0 {
+		for i, e := range res.Errors {
+			if i == 0 {
+				r.Error = e.Message
+			} else {
+				r.Error = fmt.Sprintf("%v, %v", r.Error, e.Message)
+			}
+		}
+	}
+
+	return nil
+}

lib/results_test.go

@@ -96,10 +96,9 @@ func TestResultEncoding(t *testing.T) {
 					BytesIn:   rapid.Uint64().Draw(t, "bytes_in"),
 					BytesOut:  rapid.Uint64().Draw(t, "bytes_out"),
 					Error:     rapid.StringMatching(`^\w+$`).Draw(t, "error"),
-					Body:      rapid.SliceOf(rapid.Byte()).Draw(t, "body"),
-					Method: rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$").
-						Draw(t, "method"),
-					URL: rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&%\$#\=~]*)$`).Draw(t, "url"),
+					Body:      []byte("{\"data\":{\"vegeta\":\"punch\"}}"),
+					Method:    rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$").Draw(t, "method"),
+					URL:       rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&%\$#\=~]*)$`).Draw(t, "url"),
 				}
 
 				if len(hdrs) > 0 {
@@ -115,7 +114,7 @@ func TestResultEncoding(t *testing.T) {
 				var buf bytes.Buffer
 				enc := tc.enc(&buf)
 				for j := 0; j < 2; j++ {
-					if err := enc(&want); err != nil {
+					if err := enc.Encode(&want); err != nil {
 						t.Fatal(err)
 					}
 				}
@@ -128,7 +127,7 @@ func TestResultEncoding(t *testing.T) {
 				}
 				for j := 0; j < 2; j++ {
 					var got Result
-					if err := dec(&got); err != nil {
+					if err := dec.Decode(&got); err != nil {
 						t.Fatalf("err: %q buffer: %s", err, encoded)
 					}
 
@@ -142,6 +141,67 @@ func TestResultEncoding(t *testing.T) {
 	}
 }
 
+func TestGQLEncoding(t *testing.T) {
+	rapid.Check(t, func(t *rapid.T) {
+		hdrs := rapid.MapOf(
+			rapid.StringMatching("^[!#$%&'*+\\-.^_`|~0-9a-zA-Z]+$"),
+			rapid.SliceOfN(rapid.StringMatching(`^[0-9a-zA-Z]+$`), 1, -1),
+		).Draw(t, "headers")
+
+		in := Result{
+			Attack:    "gql-test",
+			Seq:       rapid.Uint64().Draw(t, "seq"),
+			Code:      200,
+			Timestamp: time.Unix(rapid.Int64Range(0, 1e8).Draw(t, "timestamp"), 0),
+			Latency:   time.Duration(rapid.Int64Min(0).Draw(t, "latency")),
+			BytesIn:   rapid.Uint64().Draw(t, "bytes_in"),
+			BytesOut:  rapid.Uint64().Draw(t, "bytes_out"),
+			Error:     "",
+			Body:      []byte("{\"data\":{},\"errors\":[{\"message\":\"no punch\"}]}"),
+			Method:    rapid.StringMatching("^(GET|PUT|POST|DELETE|HEAD|OPTIONS)$").Draw(t, "method"),
+			URL:       rapid.StringMatching(`^(https?):\/\/([a-zA-Z0-9-\.]+)(:[0-9]{1,5})?\/?([a-zA-Z0-9\-\._\?\,\'\/\\\+&amp;%\$#\=~]*)$`).Draw(t, "url"),
+		}
+
+		if len(hdrs) > 0 {
+			in.Headers = make(http.Header, len(hdrs))
+		}
+
+		for k, vs := range hdrs {
+			for _, v := range vs {
+				in.Headers.Add(k, v)
+			}
+		}
+
+		want := in
+		want.Error = "no punch"
+
+		var buf bytes.Buffer
+		enc := NewJSONEncoder(&buf)
+		for j := 0; j < 2; j++ {
+			if err := enc.Encode(&in); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		encoded := buf.String()
+
+		dec := NewJSONDecoder(&buf)
+		for j := 0; j < 2; j++ {
+			var got Result
+			if err := dec.Decode(&got); err != nil {
+				t.Fatalf("err: %q buffer: %s", err, encoded)
+			}
+
+			AsGraphQL(&got)
+
+			if !got.Equal(want) {
+				t.Logf("encoded: %s", encoded)
+				t.Fatalf("mismatch: %s", cmp.Diff(got, want))
+			}
+		}
+	})
+}
+
 func BenchmarkResultEncodings(b *testing.B) {
 	b.StopTimer()
 	b.ResetTimer()