From c61d011076f708caa4f76cf2dc9f951a8720e2a8 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Tue, 11 Jul 2023 14:45:50 +0100
Subject: [PATCH] Support RestClient in MockRestServiceServer

Closes gh-30821
---
 .../web/client/MockRestServiceServer.java     | 67 ++++++++++++++-----
 1 file changed, 52 insertions(+), 15 deletions(-)

diff --git a/spring-test/src/main/java/org/springframework/test/web/client/MockRestServiceServer.java b/spring-test/src/main/java/org/springframework/test/web/client/MockRestServiceServer.java
index d51333d8629b..c67f6ba44496 100644
--- a/spring-test/src/main/java/org/springframework/test/web/client/MockRestServiceServer.java
+++ b/spring-test/src/main/java/org/springframework/test/web/client/MockRestServiceServer.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,7 @@
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.mock.http.client.MockClientHttpRequest;
 import org.springframework.util.Assert;
+import org.springframework.web.client.RestClient;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.client.support.RestGatewaySupport;
 
@@ -137,13 +138,22 @@ public void reset() {
 	}
 
 
+	/**
+	 * Return a builder for a {@code MockRestServiceServer} that should be used
+	 * to reply to the given {@code RestTemplate}.
+	 * @since 6.1
+	 */
+	public static MockRestServiceServerBuilder bindTo(RestClient.Builder restClientBuilder) {
+		return new RestClientMockRestServiceServerBuilder(restClientBuilder);
+	}
+
 	/**
 	 * Return a builder for a {@code MockRestServiceServer} that should be used
 	 * to reply to the given {@code RestTemplate}.
 	 * @since 4.3
 	 */
 	public static MockRestServiceServerBuilder bindTo(RestTemplate restTemplate) {
-		return new DefaultBuilder(restTemplate);
+		return new RestTemplateMockRestServiceServerBuilder(restTemplate);
 	}
 
 	/**
@@ -153,7 +163,7 @@ public static MockRestServiceServerBuilder bindTo(RestTemplate restTemplate) {
 	 */
 	public static MockRestServiceServerBuilder bindTo(RestGatewaySupport restGatewaySupport) {
 		Assert.notNull(restGatewaySupport, "'restGatewaySupport' must not be null");
-		return new DefaultBuilder(restGatewaySupport.getRestTemplate());
+		return new RestTemplateMockRestServiceServerBuilder(restGatewaySupport.getRestTemplate());
 	}
 
 
@@ -214,20 +224,14 @@ public interface MockRestServiceServerBuilder {
 	}
 
 
-	private static class DefaultBuilder implements MockRestServiceServerBuilder {
 
-		private final RestTemplate restTemplate;
+	private abstract static class AbstractMockRestServiceServerBuilder implements MockRestServiceServerBuilder {
 
 		private boolean ignoreExpectOrder;
 
 		private boolean bufferContent;
 
 
-		public DefaultBuilder(RestTemplate restTemplate) {
-			Assert.notNull(restTemplate, "RestTemplate must not be null");
-			this.restTemplate = restTemplate;
-		}
-
 		@Override
 		public MockRestServiceServerBuilder ignoreExpectOrder(boolean ignoreExpectOrder) {
 			this.ignoreExpectOrder = ignoreExpectOrder;
@@ -253,15 +257,48 @@ public MockRestServiceServer build() {
 		@Override
 		public MockRestServiceServer build(RequestExpectationManager manager) {
 			MockRestServiceServer server = new MockRestServiceServer(manager);
-			MockClientHttpRequestFactory factory = server.new MockClientHttpRequestFactory();
+			ClientHttpRequestFactory factory = server.new MockClientHttpRequestFactory();
 			if (this.bufferContent) {
-				this.restTemplate.setRequestFactory(new BufferingClientHttpRequestFactory(factory));
-			}
-			else {
-				this.restTemplate.setRequestFactory(factory);
+				factory = new BufferingClientHttpRequestFactory(factory);
 			}
+			injectRequestFactory(factory);
 			return server;
 		}
+
+		protected abstract void injectRequestFactory(ClientHttpRequestFactory requestFactory);
+
+	}
+
+
+	private static class RestClientMockRestServiceServerBuilder extends AbstractMockRestServiceServerBuilder {
+
+		private final RestClient.Builder restClientBuilder;
+
+		RestClientMockRestServiceServerBuilder(RestClient.Builder restClientBuilder) {
+			Assert.notNull(restClientBuilder, "RestClient.Builder must not be null");
+			this.restClientBuilder = restClientBuilder;
+		}
+
+		@Override
+		protected void injectRequestFactory(ClientHttpRequestFactory requestFactory) {
+			this.restClientBuilder.requestFactory(requestFactory);
+		}
+	}
+
+
+	private static class RestTemplateMockRestServiceServerBuilder extends AbstractMockRestServiceServerBuilder {
+
+		private final RestTemplate restTemplate;
+
+		RestTemplateMockRestServiceServerBuilder(RestTemplate restTemplate) {
+			Assert.notNull(restTemplate, "RestTemplate must not be null");
+			this.restTemplate = restTemplate;
+		}
+
+		@Override
+		protected void injectRequestFactory(ClientHttpRequestFactory requestFactory) {
+			this.restTemplate.setRequestFactory(requestFactory);
+		}
 	}