From a809d7c79d95e33c5769df627aa4c789f9909945 Mon Sep 17 00:00:00 2001
From: Stein Fletcher <steindfletcher@gmail.com>
Date: Mon, 1 Mar 2021 14:55:25 +0000
Subject: [PATCH] Support JSON marshalling any type when setting the mock
 req/res body

This aligns the mock API with the apitest API for creating a request
---
 mocks.go      | 28 ++++++++++++++++++++++++++++
 mocks_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/mocks.go b/mocks.go
index a468aed..0deda0e 100644
--- a/mocks.go
+++ b/mocks.go
@@ -487,6 +487,20 @@ func (r *MockRequest) BodyFromFile(f string) *MockRequest {
 	return r
 }
 
+// JSON is a convenience method for setting the mock request body
+func (r *MockRequest) JSON(v interface{}) *MockRequest {
+	switch x := v.(type) {
+	case string:
+		r.body = x
+	case []byte:
+		r.body = string(x)
+	default:
+		asJSON, _ := json.Marshal(x)
+		r.body = string(asJSON)
+	}
+	return r
+}
+
 // Header configures the mock request to match the given header
 func (r *MockRequest) Header(key, value string) *MockRequest {
 	normalizedKey := textproto.CanonicalMIMEHeaderKey(key)
@@ -660,6 +674,20 @@ func (r *MockResponse) BodyFromFile(f string) *MockResponse {
 	return r
 }
 
+// JSON is a convenience method for setting the mock response body
+func (r *MockResponse) JSON(v interface{}) *MockResponse {
+	switch x := v.(type) {
+	case string:
+		r.body = x
+	case []byte:
+		r.body = string(x)
+	default:
+		asJSON, _ := json.Marshal(x)
+		r.body = string(asJSON)
+	}
+	return r
+}
+
 // Status respond with the given status
 func (r *MockResponse) Status(statusCode int) *MockResponse {
 	r.statusCode = statusCode
diff --git a/mocks_test.go b/mocks_test.go
index ce30b40..f6da78d 100644
--- a/mocks_test.go
+++ b/mocks_test.go
@@ -635,6 +635,23 @@ func TestMocks_BodyMatcher(t *testing.T) {
 	}
 }
 
+func TestMocks_RequestBody(t *testing.T) {
+	tests := map[string]struct {
+		requestBody interface{}
+	}{
+		"supports string input": {`{"a":1}`},
+		"supports maps":         {map[string]int{"a": 1}},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "/path", strings.NewReader(`{"a":1}`))
+			err := bodyMatcher(req, NewMock().Get("/").JSON(test.requestBody))
+			assert.NoError(t, err)
+		})
+	}
+}
+
 func TestMocks_PathMatcher(t *testing.T) {
 	tests := []struct {
 		requestUrl    string
@@ -952,6 +969,19 @@ func TestMocks_Response_SetsTheBodyAsJSON(t *testing.T) {
 	assert.Equal(t, "application/json", response.Header.Get("Content-Type"))
 }
 
+func TestMocks_ResponseJSON(t *testing.T) {
+	mockResponse := NewMock().
+		Get("assert").
+		RespondWith().
+		JSON(map[string]int{"a": 123})
+
+	response := buildResponseFromMock(mockResponse)
+
+	bytes, _ := ioutil.ReadAll(response.Body)
+	assert.Equal(t, string(bytes), `{"a":123}`)
+	assert.Equal(t, "application/json", response.Header.Get("Content-Type"))
+}
+
 func TestMocks_Response_SetsTheBodyAsOther(t *testing.T) {
 	mockResponse := NewMock().
 		Get("assert").