From 2fb12a35b9c216ee6e916ee35886c03f44f63111 Mon Sep 17 00:00:00 2001
From: Ryan Bergman <Ryan.Bergman@gmail.com>
Date: Tue, 12 Sep 2023 20:26:31 -0500
Subject: [PATCH] addresses issue #493 make the mock client work with the mock
 response because that is what people want

---
 .../kong/unirest/core/ExpectedResponse.java   | 27 +++++++++++++++++++
 .../java/kong/unirest/core/Invocation.java    | 18 ++++++++++++-
 .../src/test/java/kong/tests/AssertTest.java  | 26 ++++++++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/unirest-mocks/src/main/java/kong/unirest/core/ExpectedResponse.java b/unirest-mocks/src/main/java/kong/unirest/core/ExpectedResponse.java
index 3f1ff654..54713536 100644
--- a/unirest-mocks/src/main/java/kong/unirest/core/ExpectedResponse.java
+++ b/unirest-mocks/src/main/java/kong/unirest/core/ExpectedResponse.java
@@ -28,9 +28,36 @@
 import kong.unirest.core.json.JSONElement;
 
 public interface ExpectedResponse {
+    /**
+     * adds a header to the expected response
+     * @param key the header key
+     * @param key the header value
+     * @return this ExpectedResponse
+     */
     ExpectedResponse withHeader(String key, String value);
+
+    /**
+     * adds a collection of headers to the expected response
+     * @param headers the headers
+     * @return This ExpectedResponse
+     */
+    ExpectedResponse withHeaders(Headers headers);
+
+    /**
+     * sets the status of the expected response
+     * @param httpStatus the http status code
+     * @return this ExpectedResponse
+     */
     ExpectedResponse withStatus(int httpStatus);
+
+    /**
+     * sets the status of the expected response
+     * @param httpStatus the http status code
+     * @param statusMessage the status message
+     * @return this ExpectedResponse
+     */
     ExpectedResponse withStatus(int httpStatus, String statusMessage);
+
     /**
      * expect a string response
      * @param body the expected response body
diff --git a/unirest-mocks/src/main/java/kong/unirest/core/Invocation.java b/unirest-mocks/src/main/java/kong/unirest/core/Invocation.java
index baf25175..0166b022 100644
--- a/unirest-mocks/src/main/java/kong/unirest/core/Invocation.java
+++ b/unirest-mocks/src/main/java/kong/unirest/core/Invocation.java
@@ -84,10 +84,20 @@ public ExpectedResponse thenReturn(JSONElement jsonObject) {
 
     @Override
     public ExpectedResponse thenReturn(Object pojo) {
-        this.response = o -> o.writeValue(pojo);
+        if(pojo instanceof MockResponse){
+            var res = (MockResponse)pojo;
+            return thenReturn(res);
+        } else {
+            this.response = o -> o.writeValue(pojo);
+        }
         return this;
     }
 
+    private ExpectedResponse thenReturn(MockResponse res) {
+        this.response = o -> res.getBody() == null ? null : String.valueOf(res.getBody());
+        return withStatus(res.getStatus()).withHeaders(res.getHeaders());
+    }
+
     @Override
     public ExpectedResponse thenReturn(Supplier<String> supplier) {
         this.response = o -> supplier.get();
@@ -228,6 +238,12 @@ public ExpectedResponse withHeader(String key, String value) {
         return this;
     }
 
+    @Override
+    public ExpectedResponse withHeaders(Headers value) {
+        value.all().forEach(h -> withHeader(h.getName(), h.getValue()));
+        return this;
+    }
+
     @Override
     public ExpectedResponse withStatus(int httpStatus) {
         return withStatus(httpStatus,"");
diff --git a/unirest-mocks/src/test/java/kong/tests/AssertTest.java b/unirest-mocks/src/test/java/kong/tests/AssertTest.java
index ef7d8e84..562616f8 100644
--- a/unirest-mocks/src/test/java/kong/tests/AssertTest.java
+++ b/unirest-mocks/src/test/java/kong/tests/AssertTest.java
@@ -28,6 +28,7 @@
 import kong.unirest.core.*;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
 import java.util.function.Supplier;
 
 import static kong.unirest.core.HttpMethod.GET;
@@ -223,6 +224,31 @@ void assertPostWithNoBody() {
                 "\t null", ex.getMessage());
     }
 
+    @Test
+    void returnAMockResponseObject() {
+        client.expect(HttpMethod.POST, path)
+                .thenReturn(MockResponse.of(500, "error")
+                        .withHeader("cool", "beans"));
+
+        HttpResponse<String> response = Unirest.post(path).asString();
+
+        assertEquals(500, response.getStatus());
+        assertEquals("error", response.getBody());
+        assertEquals(List.of("beans"), response.getHeaders().get("cool"));
+    }
+
+    @Test
+    void returnAMockResponseObjectWithoutStuff() {
+        client.expect(HttpMethod.POST, path)
+                .thenReturn(MockResponse.of(500, null));
+
+        HttpResponse<String> response = Unirest.post(path).asString();
+
+        assertEquals(500, response.getStatus());
+        assertEquals(null, response.getBody());
+        assertEquals(List.of(), response.getHeaders().get("cool"));
+    }
+
     private static class BodyBuddy implements Supplier<String>{
         String body;
         @Override