From aae13fd6d56ac42a869f33b90eb1fc2d2fb33e25 Mon Sep 17 00:00:00 2001
From: Yann D'Isanto <le-yams@users.noreply.github.com>
Date: Tue, 15 Oct 2024 10:03:19 +0200
Subject: [PATCH] feat: stub with bytes (#6)

---
 stubs.go      | 18 +++++++++++-------
 stubs_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/stubs.go b/stubs.go
index b3bf3da..8c82ac4 100644
--- a/stubs.go
+++ b/stubs.go
@@ -23,15 +23,19 @@ func (stub *StubBuilder) WithStatusCode(statusCode int) *APIMock {
 }
 
 func (stub *StubBuilder) WithJSON(statusCode int, content interface{}) *APIMock {
+	body, err := json.Marshal(content)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return stub.WithBody(statusCode, body, "application/json")
+}
+
+func (stub *StubBuilder) WithBody(statusCode int, body []byte, contentType string) *APIMock {
 	return stub.With(func(writer http.ResponseWriter, request *http.Request) {
-		writer.Header().Add("Content-Type", "application/json")
+		writer.Header().Add("Content-Type", contentType)
 		writer.WriteHeader(statusCode)
-
-		bytes, err := json.Marshal(content)
-		if err != nil {
-			log.Fatal(err)
-		}
-		_, err = writer.Write(bytes)
+		_, err := writer.Write(body)
 		if err != nil {
 			log.Fatal(err)
 		}
diff --git a/stubs_test.go b/stubs_test.go
index 08da08e..26cbefb 100644
--- a/stubs_test.go
+++ b/stubs_test.go
@@ -75,6 +75,10 @@ func TestApiStubbedEndpointWithJson(t *testing.T) {
 	// Assert
 	testState.assertDidNotFailed()
 
+	e.GET("/endpoint").
+		Expect().
+		Header("Content-Type").IsEqual("application/json")
+
 	responseObject := e.GET("/endpoint").
 		Expect().
 		Status(http.StatusOK).
@@ -82,3 +86,31 @@ func TestApiStubbedEndpointWithJson(t *testing.T) {
 
 	responseObject.Value("value").IsEqual("Hello")
 }
+
+func TestApiStubbedEndpointWithBody(t *testing.T) {
+    t.Parallel()
+	// Arrange
+	testState := NewTestingMock(t)
+	mockedAPI := API(testState)
+	defer func() { mockedAPI.Close() }()
+	body := []byte("Hello!")
+
+	mockedAPI.
+		Stub(http.MethodGet, "/endpoint").
+		WithBody(http.StatusOK, body, "text/plain")
+
+	// Act
+	e := httpexpect.Default(t, mockedAPI.GetURL().String())
+
+	// Assert
+	testState.assertDidNotFailed()
+
+	e.GET("/endpoint").
+		Expect().
+		Header("Content-Type").IsEqual("text/plain")
+
+	e.GET("/endpoint").
+		Expect().
+		Status(http.StatusOK).
+		Body().IsEqual("Hello!")
+}