Capture response body in Results

Fixes #279
tsenart · May 18, 2018 · 2da25b2 · 2da25b2
1 parent da7a2c6
commit 2da25b2
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 6 deletions.
diff --git a/lib/attack.go b/lib/attack.go
@@ -3,7 +3,6 @@ package vegeta
 import (
 	"crypto/tls"
 	"fmt"
-	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -263,11 +262,10 @@ func (a *Attacker) hit(tr Targeter, tm time.Time) *Result {
 	}
 	defer r.Body.Close()
 
-	in, err := io.Copy(ioutil.Discard, r.Body)
-	if err != nil {
+	if res.Body, err = ioutil.ReadAll(r.Body); err != nil {
 		return &res
 	}
-	res.BytesIn = uint64(in)
+	res.BytesIn = uint64(len(res.Body))
 
 	if req.ContentLength != -1 {
 		res.BytesOut = uint64(req.ContentLength)

diff --git a/lib/attack_test.go b/lib/attack_test.go
@@ -1,6 +1,7 @@
 package vegeta
 
 import (
+	"bytes"
 	"crypto/tls"
 	"fmt"
 	"io"
@@ -184,3 +185,21 @@ func TestBadTargeterError(t *testing.T) {
 		t.Fatalf("got: %v, want: %v", got, want)
 	}
 }
+
+func TestResponseBodyCapture(t *testing.T) {
+	t.Parallel()
+
+	want := []byte("VEGETA")
+	server := httptest.NewServer(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write(want)
+		}),
+	)
+	defer server.Close()
+	atk := NewAttacker()
+	tr := NewStaticTargeter(Target{Method: "GET", URL: server.URL})
+	res := atk.hit(tr, time.Now())
+	if got := res.Body; !bytes.Equal(got, want) {
+		t.Fatalf("got: %v, want: %v", got, want)
+	}
+}
diff --git a/lib/results.go b/lib/results.go
@@ -1,6 +1,8 @@
 package vegeta
 
 import (
+	"bytes"
+	"encoding/base64"
 	"encoding/csv"
 	"encoding/gob"
 	"encoding/json"
@@ -22,6 +24,7 @@ type Result struct {
 	BytesOut  uint64        `json:"bytes_out"`
 	BytesIn   uint64        `json:"bytes_in"`
 	Error     string        `json:"error"`
+	Body      []byte        `json:"body"`
 }
 
 // End returns the time at which a Result ended.
@@ -34,7 +37,8 @@ func (r Result) Equal(other Result) bool {
 		r.Latency == other.Latency &&
 		r.BytesIn == other.BytesIn &&
 		r.BytesOut == other.BytesOut &&
-		r.Error == other.Error
+		r.Error == other.Error &&
+		bytes.Equal(r.Body, other.Body)
 }
 
 // Results is a slice of Result type elements.
@@ -110,6 +114,7 @@ func NewCSVEncoder(w io.Writer) Encoder {
 			strconv.FormatUint(r.BytesOut, 10),
 			strconv.FormatUint(r.BytesIn, 10),
 			r.Error,
+			base64.StdEncoding.EncodeToString(r.Body),
 		})
 
 		if err != nil {
@@ -158,6 +163,7 @@ func NewCSVDecoder(rd io.Reader) Decoder {
 		}
 
 		r.Error = rec[5]
+		r.Body, err = base64.StdEncoding.DecodeString(rec[6])
 
 		return err
 	}

diff --git a/lib/results_test.go b/lib/results_test.go
@@ -63,14 +63,15 @@ func TestEncoding(t *testing.T) {
 			var buf bytes.Buffer
 			enc := tc.enc(&buf)
 			dec := tc.dec(&buf)
-			err := quick.Check(func(code uint16, ts uint32, latency time.Duration, bsIn, bsOut uint64, e string) bool {
+			err := quick.Check(func(code uint16, ts uint32, latency time.Duration, bsIn, bsOut uint64, body []byte, e string) bool {
 				want := Result{
 					Code:      code,
 					Timestamp: time.Unix(int64(ts), 0),
 					Latency:   latency,
 					BytesIn:   bsIn,
 					BytesOut:  bsOut,
 					Error:     e,
+					Body:      body,
 				}
 
 				if err := enc(&want); err != nil {