better tests/separation for response handling

responses/responses.go

@@ -1,11 +1,13 @@
 package responses
 
 import (
-	"bytes"
 	"crypto/md5"
 	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
+	"github.com/tednaleid/ganda/config"
 	"github.com/tednaleid/ganda/execcontext"
+	"hash"
 	"io"
 	"net/http"
 	"os"
@@ -19,13 +21,26 @@ func StartResponseWorkers(responses <-chan *http.Response, context *execcontext.
 
 	for i := 1; i <= context.ResponseWorkers; i++ {
 		go func() {
+			/*
+				high level algorithm
+					- if we're saving files, we want each response to create it's own file and have an io.Writer
+					- if we're emitting to stdout, we've already got an io.Writer at context.Out
+					- this needs to be called per response, so we can't just pass in the io.Writer to the worker
+					- if we're emitting to stdout, we want to emit a \n between each non zero length response
+				if we're emitting a json-envelope
+					-  we want to emit part of the JSON up to the body, then emit the body function, then close the JSON
+					- if we're not emitting JSON envelope, no need to wrap it
+				separately, we want the response-body to take an io.writer from the above and be able to write to it
+					- create higher-order function that returns a function that takes a response and an io.Writer
+			*/
+
 			emitResponse := determineEmitResponseFn(context)
 
 			if context.WriteFiles {
 				// TODO change this so that it calls determineEmitResponseFn regardless of if it is saving or printing
 				// and then passes that responseFn to the worker we create
 
-				responseSavingWorker(responses, context)
+				responseSavingWorker(responses, emitResponse, context)
 			} else {
 				responsePrintingWorker(responses, emitResponse, context)
 			}
@@ -36,11 +51,13 @@ func StartResponseWorkers(responses <-chan *http.Response, context *execcontext.
 	return &responseWaitGroup
 }
 
-func responseSavingWorker(responses <-chan *http.Response, context *execcontext.Context) {
+func responseSavingWorker(responses <-chan *http.Response, emitResponse emitResponseFn, context *execcontext.Context) {
 	specialCharactersRegexp := regexp.MustCompile("[^A-Za-z0-9]+")
 
 	responseWorker(responses, func(response *http.Response) {
 		filename := specialCharactersRegexp.ReplaceAllString(response.Request.URL.String(), "-")
+		// this should be changed to get an io.Writer for the file that we can pass in
+
 		fullPath := saveBodyToFile(context.BaseDirectory, context.SubdirLength, filename, response.Body)
 		context.Logger.LogResponse(response.StatusCode, response.Request.URL.String()+" -> "+fullPath)
 	})
@@ -54,75 +71,126 @@ func responsePrintingWorker(responses <-chan *http.Response, emitResponse emitRe
 	})
 }
 
-type emitResponseFn func(response *http.Response, out io.Writer)
+// takes a response and writes it to the writer, returns true if it wrote anything
+type emitResponseFn func(response *http.Response, out io.Writer) (written int64, err error)
 
+// we might wrap the body response in a JSON envelope
 func determineEmitResponseFn(context *execcontext.Context) emitResponseFn {
+	bodyResponseFn := determineEmitBodyResponseFn(context)
+
 	if context.JsonEnvelope {
-		if context.DiscardBody {
-			return emitJsonMessagesWithoutBody
-		} else if context.HashBody {
-			return emitJsonMessageSha256
-		}
-		return emitJsonMessages
+		// it matters what the body response function is:
+		// if it's raw, we want to just emit it
+		// if it's discard, we want to emit `null`. still need to call emitNothing so we close the body
+		// if it's base64, we want to encapsulate it in quotes for a JSON string
+		// if it's escaped, we want to encapsulate it in quotes for a JSON string
+		// if it's sha256, we want to encapsulate it in quotes for a JSON string
+
+		//if context.DiscardBody {
+		//	return emitJsonMessagesWithoutBody
+		//} else if context.HashBody {
+		//	return emitJsonMessageSha256
+		//}
+		//return emitJsonMessages
 	}
 
-	if context.DiscardBody {
+	return bodyResponseFn
+}
+
+func determineEmitBodyResponseFn(context *execcontext.Context) emitResponseFn {
+	switch context.ResponseBody {
+	case config.Raw:
+		return emitRawBody
+	case config.Sha256:
+		return emitSha256BodyFn()
+	case config.Discard:
 		return emitNothing
+	case config.Escaped:
+		return emitNothing //TODO
+	case config.Base64:
+		return emitNothing //TODO
+	default:
+		panic(fmt.Sprintf("unknown response body type %s", context.ResponseBody))
 	}
-
-	return emitRawMessages
 }
 
-func emitRawMessages(response *http.Response, out io.Writer) {
+func emitRawBody(response *http.Response, out io.Writer) (written int64, err error) {
 	defer response.Body.Close()
-	buf := new(bytes.Buffer)
-	buf.ReadFrom(response.Body)
-
-	if buf.Len() > 0 {
-		buf.WriteByte('\n')
-		out.Write(buf.Bytes())
-	}
+	return io.Copy(out, response.Body)
 }
 
-func emitNothing(response *http.Response, out io.Writer) {
-	response.Body.Close()
+func emitSha256BodyFn() func(response *http.Response, out io.Writer) (written int64, err error) {
+	hasher := sha256.New()
+	return func(response *http.Response, out io.Writer) (written int64, err error) {
+		return emitHashedBody(hasher, response, out)
+	}
 }
 
-func emitJsonMessages(response *http.Response, out io.Writer) {
+func emitHashedBody(hasher hash.Hash, response *http.Response, out io.Writer) (written int64, err error) {
 	defer response.Body.Close()
-	buf := new(bytes.Buffer)
-	buf.ReadFrom(response.Body)
 
-	if buf.Len() > 0 {
-		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": %s }\n", response.Request.URL.String(), response.StatusCode, buf.Len(), buf)
-	} else {
-		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": null }\n", response.Request.URL.String(), response.StatusCode, 0)
+	hasher.Reset()
+	if _, err := io.Copy(hasher, response.Body); err != nil {
+		return 0, err
 	}
+
+	n, err := fmt.Fprint(out, hex.EncodeToString(hasher.Sum(nil)))
+	return int64(n), err
 }
 
-func emitJsonMessagesWithoutBody(response *http.Response, out io.Writer) {
+func emitNothing(response *http.Response, out io.Writer) (written int64, err error) {
 	response.Body.Close()
-	fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d}\n", response.Request.URL.String(), response.StatusCode)
+	return 0, nil
 }
 
-func emitJsonMessageSha256(response *http.Response, out io.Writer) {
-	defer response.Body.Close()
-	buf := new(bytes.Buffer)
-	buf.ReadFrom(response.Body)
-
-	if buf.Len() > 0 {
-		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": \"%x\" }\n", response.Request.URL.String(), response.StatusCode, buf.Len(), sha256.Sum256(buf.Bytes()))
-	} else {
-		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": null }\n", response.Request.URL.String(), response.StatusCode, 0)
-	}
-}
+//func emitJsonMessages(response *http.Response, out io.Writer) {
+//	defer response.Body.Close()
+//	buf := new(bytes.Buffer)
+//	buf.ReadFrom(response.Body)
+//
+//	if buf.Len() > 0 {
+//		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": %s }\n", response.Request.URL.String(), response.StatusCode, buf.Len(), buf)
+//	} else {
+//		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": null }\n", response.Request.URL.String(), response.StatusCode, 0)
+//	}
+//}
+
+//func emitJsonMessagesWithoutBody(response *http.Response, out io.Writer) (written int64, err error) {
+//	response.Body.Close()
+//	return fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d}\n", response.Request.URL.String(), response.StatusCode)
+//}
+//
+//func emitJsonMessageSha256(response *http.Response, out io.Writer) bool {
+//	defer response.Body.Close()
+//	buf := new(bytes.Buffer)
+//	buf.ReadFrom(response.Body)
+//
+//	if buf.Len() > 0 {
+//		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": \"%x\" }\n", response.Request.URL.String(), response.StatusCode, buf.Len(), sha256.Sum256(buf.Bytes()))
+//	} else {
+//		fmt.Fprintf(out, "{ \"url\": \"%s\", \"code\": %d, \"length\": %d, \"body\": null }\n", response.Request.URL.String(), response.StatusCode, 0)
+//	}
+//	return true
+//}
 
 func responseWorker(responses <-chan *http.Response, responseHandler func(*http.Response)) {
 	for response := range responses {
 		responseHandler(response)
 	}
 }
 
+func createWritableFile(baseDirectory string, subdirLength int64, filename string) io.WriteCloser {
+	directory := directoryForFile(baseDirectory, filename, subdirLength)
+	fullPath := directory + filename
+
+	file, err := os.Create(fullPath)
+	if err != nil {
+		panic(err)
+	}
+
+	return file
+}
+
 func saveBodyToFile(baseDirectory string, subdirLength int64, filename string, body io.ReadCloser) string {
 	defer body.Close()
 
@@ -133,6 +201,7 @@ func saveBodyToFile(baseDirectory string, subdirLength int64, filename string, b
 	if err != nil {
 		panic(err)
 	}
+	defer file.Close()
 
 	_, err = io.Copy(file, body)
 	if err != nil {

responses/responses_test.go

@@ -1 +1,93 @@
 package responses
+
+import (
+	"bytes"
+	"github.com/stretchr/testify/assert"
+	"github.com/tednaleid/ganda/config"
+	"github.com/tednaleid/ganda/execcontext"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func TestRawOutput(t *testing.T) {
+	context := &execcontext.Context{
+		ResponseBody: config.Raw,
+	}
+
+	responseFn := determineEmitResponseFn(context)
+	assert.NotNil(t, responseFn)
+
+	response, responseBody := mockResponseBody("hello world")
+	writeCloser := NewMockWriteCloser()
+
+	responseFn(response, writeCloser)
+
+	assert.True(t, responseBody.Closed)
+	assert.Equal(t, "hello world", writeCloser.ToString())
+}
+
+func TestSha256Output(t *testing.T) {
+	context := &execcontext.Context{
+		ResponseBody: config.Sha256,
+	}
+
+	responseFn := determineEmitResponseFn(context)
+	assert.NotNil(t, responseFn)
+
+	response, responseBody := mockResponseBody("hello world")
+	out := NewMockWriteCloser()
+
+	responseFn(response, out)
+
+	assert.True(t, responseBody.Closed)
+	assert.Equal(t, "a948904f2f0f479b8f8197694b30184b0d2ed1c1cd2a1ec0fb85d299a192a447", out.ToString())
+}
+
+func mockResponseBody(body string) (*http.Response, *MockReadCloser) {
+	responseBody := &MockReadCloser{ReadCloser: io.NopCloser(strings.NewReader(body)), Closed: false}
+	return &http.Response{
+		StatusCode: 200,
+		Body:       responseBody,
+	}, responseBody
+}
+
+type MockReadCloser struct {
+	ReadCloser io.ReadCloser
+	Closed     bool
+}
+
+func (m *MockReadCloser) Read(p []byte) (n int, err error) {
+	return m.ReadCloser.Read(p)
+}
+
+func (m *MockReadCloser) Close() error {
+	m.Closed = true
+	return nil
+}
+
+type MockWriteCloser struct {
+	Buffer *bytes.Buffer
+	Closed bool
+}
+
+func (m *MockWriteCloser) Write(p []byte) (n int, err error) {
+	return m.Buffer.Write(p)
+}
+
+func (m *MockWriteCloser) Close() error {
+	m.Closed = true
+	return nil
+}
+
+func (m *MockWriteCloser) ToString() string {
+	return m.Buffer.String()
+}
+
+func NewMockWriteCloser() *MockWriteCloser {
+	return &MockWriteCloser{
+		Buffer: new(bytes.Buffer),
+		Closed: false,
+	}
+}