Improve HttpHandlerConnection completion

Before this commit the connector waited for a completed response (via
ServerHttpResponse#setComplete or ServerHttpResponse#writeWith) or an
error signal in handling, but it didn't deal explicitly with the case
where both can occur.

This commit explicitly waits for the completion of handling (success
or error) before passing the response downstream. If an error occurs
after response completion, it is wrapped in a dedicated exception that
also provides access to the completed response.

Close gh-24051

spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpResponse.java

@@ -148,4 +148,10 @@ private Charset getCharset() {
 		return (charset != null ? charset : StandardCharsets.UTF_8);
 	}
 
+
+	@Override
+	public String toString() {
+		HttpStatus code = HttpStatus.resolve(this.status);
+		return (code != null ? code.name() + "(" + this.status + ")" : "Status (" + this.status + ")") + this.headers;
+	}
 }

spring-test/src/main/java/org/springframework/test/web/reactive/server/HttpHandlerConnector.java

@@ -83,7 +83,9 @@ public Mono<ClientHttpResponse> connect(HttpMethod httpMethod, URI uri,
 	private Mono<ClientHttpResponse> doConnect(
 			HttpMethod httpMethod, URI uri, Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
 
-		MonoProcessor<ClientHttpResponse> result = MonoProcessor.create();
+		MonoProcessor<Void> requestWriteCompletion = MonoProcessor.create();
+		MonoProcessor<Void> handlerCompletion = MonoProcessor.create();
+		ClientHttpResponse[] savedResponse = new ClientHttpResponse[1];
 
 		MockClientHttpRequest mockClientRequest = new MockClientHttpRequest(httpMethod, uri);
 		MockServerHttpResponse mockServerResponse = new MockServerHttpResponse();
@@ -92,20 +94,26 @@ private Mono<ClientHttpResponse> doConnect(
 			log("Invoking HttpHandler for ", httpMethod, uri);
 			ServerHttpRequest mockServerRequest = adaptRequest(mockClientRequest, requestBody);
 			ServerHttpResponse responseToUse = prepareResponse(mockServerResponse, mockServerRequest);
-			this.handler.handle(mockServerRequest, responseToUse).subscribe(aVoid -> {}, result::onError);
+			this.handler.handle(mockServerRequest, responseToUse).subscribe(handlerCompletion);
 			return Mono.empty();
 		});
 
 		mockServerResponse.setWriteHandler(responseBody ->
 				Mono.fromRunnable(() -> {
 					log("Creating client response for ", httpMethod, uri);
-					result.onNext(adaptResponse(mockServerResponse, responseBody));
+					savedResponse[0] = adaptResponse(mockServerResponse, responseBody);
 				}));
 
 		log("Writing client request for ", httpMethod, uri);
-		requestCallback.apply(mockClientRequest).subscribe(aVoid -> {}, result::onError);
-
-		return result;
+		requestCallback.apply(mockClientRequest).subscribe(requestWriteCompletion);
+
+		return Mono.when(requestWriteCompletion, handlerCompletion)
+				.onErrorMap(ex -> {
+					ClientHttpResponse response = savedResponse[0];
+					return response != null ? new FailureAfterResponseCompletedException(response, ex) : ex;
+				})
+				.then(Mono.fromCallable(() -> savedResponse[0] != null ?
+						savedResponse[0] : adaptResponse(mockServerResponse, Flux.empty())));
 	}
 
 	private void log(String message, HttpMethod httpMethod, URI uri) {
@@ -135,4 +143,33 @@ private ClientHttpResponse adaptResponse(MockServerHttpResponse response, Flux<D
 		return clientResponse;
 	}
 
+
+	/**
+	 * Indicates that an error occurred after the server response was completed,
+	 * via {@link ServerHttpResponse#writeWith} or {@link ServerHttpResponse#setComplete()},
+	 * and can no longer be changed. This exception wraps the error and also
+	 * provides {@link #getCompletedResponse() access} to the response.
+	 * <p>What happens on an actual running server depends on when the server
+	 * commits the response and the error may or may not change the response.
+	 * Therefore in tests without a server the exception is wrapped and allowed
+	 * to propagate so the application is alerted.
+	 * @since 5.2.2
+	 */
+	@SuppressWarnings("serial")
+	public static final class FailureAfterResponseCompletedException extends RuntimeException {
+
+		private final ClientHttpResponse completedResponse;
+
+
+		private FailureAfterResponseCompletedException(ClientHttpResponse response, Throwable cause) {
+			super("Error occurred after response was completed: " + response, cause);
+			this.completedResponse = response;
+		}
+
+
+		public ClientHttpResponse getCompletedResponse() {
+			return this.completedResponse;
+		}
+	}
+
 }

spring-web/src/test/java/org/springframework/mock/http/client/reactive/test/MockClientHttpResponse.java

@@ -64,7 +64,7 @@ public MockClientHttpResponse(HttpStatus status) {
 	}
 
 	public MockClientHttpResponse(int status) {
-		Assert.isTrue(status >= 100 && status < 600, "Status must be between 1xx and 5xx");
+		Assert.isTrue(status > 99 && status < 1000, "Status must be between 100 and 999");
 		this.status = status;
 	}
 
@@ -148,4 +148,10 @@ private Charset getCharset() {
 		return (charset != null ? charset : StandardCharsets.UTF_8);
 	}
 
+
+	@Override
+	public String toString() {
+		HttpStatus code = HttpStatus.resolve(this.status);
+		return (code != null ? code.name() + "(" + this.status + ")" : "Status (" + this.status + ")") + this.headers;
+	}
 }