Add content buffering support to MockServerRestTemplateCustomizer

See gh-34833

spring-boot-project/spring-boot-test/src/main/java/org/springframework/boot/test/web/client/MockServerRestTemplateCustomizer.java

@@ -24,7 +24,9 @@ import java.util.function.Supplier;
 import org.springframework.beans.BeanUtils;
 import org.springframework.boot.web.client.RestTemplateBuilder;
 import org.springframework.boot.web.client.RestTemplateCustomizer;
+import org.springframework.http.client.BufferingClientHttpRequestFactory;
 import org.springframework.test.web.client.MockRestServiceServer;
+import org.springframework.test.web.client.MockRestServiceServer.MockRestServiceServerBuilder;
 import org.springframework.test.web.client.RequestExpectationManager;
 import org.springframework.test.web.client.SimpleRequestExpectationManager;
 import org.springframework.util.Assert;
@@ -49,6 +51,7 @@ import org.springframework.web.client.RestTemplate;
  *
  * @author Phillip Webb
  * @author Moritz Halbritter
+ * @author Chinmoy Chakraborty
  * @since 1.4.0
  * @see #getServer()
  * @see #getServer(RestTemplate)
@@ -63,6 +66,8 @@ public class MockServerRestTemplateCustomizer implements RestTemplateCustomizer
 
 	private boolean detectRootUri = true;
 
+	private boolean bufferContent = false;
+
 	public MockServerRestTemplateCustomizer() {
 		this(SimpleRequestExpectationManager::new);
 	}
@@ -96,13 +101,25 @@ public class MockServerRestTemplateCustomizer implements RestTemplateCustomizer
 		this.detectRootUri = detectRootUri;
 	}
 
+	/**
+	 * Use the {@link BufferingClientHttpRequestFactory} wrapper to buffer the input and
+	 * output streams, and for example, allow multiple reads of the response body.
+	 */
+	public void bufferContent(boolean bufferContent) {
+		this.bufferContent = bufferContent;
+	}
+
 	@Override
 	public void customize(RestTemplate restTemplate) {
 		RequestExpectationManager expectationManager = createExpectationManager();
 		if (this.detectRootUri) {
 			expectationManager = RootUriRequestExpectationManager.forRestTemplate(restTemplate, expectationManager);
 		}
-		MockRestServiceServer server = MockRestServiceServer.bindTo(restTemplate).build(expectationManager);
+		MockRestServiceServerBuilder serverBuilder = MockRestServiceServer.bindTo(restTemplate);
+		if (this.bufferContent) {
+			serverBuilder.bufferContent();
+		}
+		MockRestServiceServer server = serverBuilder.build(expectationManager);
 		this.expectationManagers.put(restTemplate, expectationManager);
 		this.servers.put(restTemplate, server);
 	}

spring-boot-project/spring-boot-test/src/test/java/org/springframework/boot/test/web/client/MockServerRestTemplateCustomizerTests.java

@@ -22,6 +22,8 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.web.client.RestTemplateBuilder;
+import org.springframework.http.client.BufferingClientHttpRequestFactory;
+import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.test.web.client.RequestExpectationManager;
 import org.springframework.test.web.client.SimpleRequestExpectationManager;
 import org.springframework.test.web.client.UnorderedRequestExpectationManager;
@@ -178,4 +180,21 @@ class MockServerRestTemplateCustomizerTests {
 		assertThat(this.customizer.getServer(template2)).extracting("expectationManager").isEqualTo(manager2);
 	}
 
+	@Test
+	void bufferContentShouldBeFalseByDefault() {
+		MockServerRestTemplateCustomizer customizer = new MockServerRestTemplateCustomizer();
+		RestTemplate restTemplate = new RestTemplate();
+		customizer.customize(restTemplate);
+		assertThat(restTemplate.getRequestFactory()).isInstanceOf(ClientHttpRequestFactory.class);
+	}
+
+	@Test
+	void enableBufferContent() {
+		MockServerRestTemplateCustomizer customizer = new MockServerRestTemplateCustomizer();
+		RestTemplate restTemplate = new RestTemplate();
+		customizer.bufferContent(true);
+		customizer.customize(restTemplate);
+		assertThat(restTemplate.getRequestFactory()).isInstanceOf(BufferingClientHttpRequestFactory.class);
+	}
+
 }