Update and expand GraphQlClient tests

See gh-10

spring-graphql/src/test/java/org/springframework/graphql/client/DefaultGraphQlClientTests.java

@@ -1,157 +0,0 @@
-/*
- * Copyright 2002-2022 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.
- * You may obtain a copy of the License at
- *
- *      https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.springframework.graphql.client;
-
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-
-import graphql.ExecutionResult;
-import graphql.ExecutionResultImpl;
-import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-
-import org.springframework.graphql.GraphQlRequest;
-import org.springframework.lang.Nullable;
-import org.springframework.util.Assert;
-
-import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
-
-/**
- *
- * @author Rossen Stoyanchev
- * @since 1.0.0
- */
-public class DefaultGraphQlClientTests {
-
-	@Test
-	void executeQuery() {
-		String query = "{" +
-				"  project(slug: \"spring-framework\") {" +
-				"    slug" +
-				"    name" +
-				"    repositoryUrl" +
-				"  }" +
-				"}";
-
-		Project expectedProject = new Project(
-				"spring-framework", "Spring Framework", "https://github.com/spring-projects/spring-framework");
-
-		ExecutionResultImpl result = ExecutionResultImpl.newExecutionResult()
-				.data(Collections.singletonMap("project", expectedProject.toMap()))
-				.build();
-
-		TestTransport transport = new TestTransport(result);
-
-		Project project = GraphQlClient.builder(transport).build()
-				.document(query)
-				.execute()
-				.map(spec -> spec.toEntity("project", Project.class))
-				.block();
-
-		assertThat(project).isNotNull();
-		assertThat(project.getSlug()).isEqualTo(expectedProject.getSlug());
-		assertThat(project.getName()).isEqualTo(expectedProject.getName());
-		assertThat(project.getRepositoryUrl()).isEqualTo(expectedProject.getRepositoryUrl());
-	}
-
-
-	private static class TestTransport implements GraphQlTransport {
-
-		private final Mono<ExecutionResult> response;
-
-		@Nullable
-		private GraphQlRequest savedRequest;
-
-		public TestTransport(ExecutionResult response) {
-			this(Mono.just(response));
-		}
-
-		public TestTransport(Mono<ExecutionResult> response) {
-			this.response = response;
-		}
-
-		public GraphQlRequest getSavedRequest() {
-			Assert.notNull(this.savedRequest, "No saved request");
-			return this.savedRequest;
-		}
-
-		@Override
-		public Mono<ExecutionResult> execute(GraphQlRequest request) {
-			this.savedRequest = request;
-			return this.response;
-		}
-
-		@Override
-		public Flux<ExecutionResult> executeSubscription(GraphQlRequest request) {
-			throw new UnsupportedOperationException();
-		}
-	}
-
-
-	private static class Project {
-
-		private String slug;
-
-		private String name;
-
-		private String repositoryUrl;
-
-		public Project() {
-		}
-
-		public Project(String slug, String name, String repositoryUrl) {
-			this.slug = slug;
-			this.name = name;
-			this.repositoryUrl = repositoryUrl;
-		}
-
-		public String getSlug() {
-			return this.slug;
-		}
-
-		public void setSlug(String slug) {
-			this.slug = slug;
-		}
-
-		public String getName() {
-			return this.name;
-		}
-
-		public void setName(String name) {
-			this.name = name;
-		}
-
-		public String getRepositoryUrl() {
-			return this.repositoryUrl;
-		}
-
-		public void setRepositoryUrl(String repositoryUrl) {
-			this.repositoryUrl = repositoryUrl;
-		}
-
-		public Map<String, Object> toMap() {
-			Map<String, Object> map = new HashMap<>();
-			map.put("slug", getSlug());
-			map.put("name", getName());
-			map.put("repositoryUrl", getRepositoryUrl());
-			return map;
-		}
-
-	}
-
-}

spring-graphql/src/test/java/org/springframework/graphql/client/GraphQlClientBuilderTests.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2002-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.client;
+
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.GraphQlRequest;
+import org.springframework.graphql.support.DocumentSource;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for the {@link GraphQlClient} builder.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class GraphQlClientBuilderTests extends GraphQlClientTestSupport {
+
+	private static final String DOCUMENT = "{ Query }";
+
+
+	@Test
+	void mutateDocumentSource() {
+
+		DocumentSource documentSource = name -> name.equals("name") ?
+				Mono.just(DOCUMENT) : Mono.error(new IllegalArgumentException());
+
+		setMockResponse("{}");
+
+		// Original
+		GraphQlClient.Builder<?> builder = graphQlClientBuilder().documentSource(documentSource);
+		GraphQlClient client = builder.build();
+		client.documentName("name").execute().block(TIMEOUT);
+
+		GraphQlRequest request = request();
+		assertThat(request.getDocument()).isEqualTo(DOCUMENT);
+
+		// Mutate
+		client = client.mutate().build();
+		client.documentName("name").execute().block(TIMEOUT);
+
+		assertThat(request().getDocument()).isEqualTo(DOCUMENT);
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/client/GraphQlClientTestSupport.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2002-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.client;
+
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.function.Consumer;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import graphql.ExecutionInput;
+import graphql.ExecutionResult;
+import graphql.ExecutionResultImpl;
+import graphql.GraphQLError;
+import org.mockito.ArgumentCaptor;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.GraphQlRequest;
+import org.springframework.graphql.RequestOutput;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Base class for {@link GraphQlClient} tests.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class GraphQlClientTestSupport {
+
+	protected static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+
+	private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+
+	private final ArgumentCaptor<GraphQlRequest> requestCaptor = ArgumentCaptor.forClass(GraphQlRequest.class);
+
+	private final GraphQlTransport transport = mock(GraphQlTransport.class);
+
+	private final GraphQlClient.Builder<?> graphQlClientBuilder = GraphQlClient.builder(this.transport);
+
+	private final GraphQlClient graphQlClient = this.graphQlClientBuilder.build();
+
+
+	protected GraphQlClient graphQlClient() {
+		return this.graphQlClient;
+	}
+
+	public GraphQlClient.Builder<?> graphQlClientBuilder() {
+		return this.graphQlClientBuilder;
+	}
+
+	protected GraphQlRequest request() {
+		return this.requestCaptor.getValue();
+	}
+
+
+	protected void setMockResponse(String data) {
+		setMockResponse(builder -> serialize(data, builder));
+	}
+
+	protected void setMockResponse(GraphQLError... errors) {
+		setMockResponse(builder -> builder.errors(Arrays.asList(errors)));
+	}
+
+	private void setMockResponse(Consumer<ExecutionResultImpl.Builder> consumer) {
+
+		ExecutionResultImpl.Builder builder = new ExecutionResultImpl.Builder();
+		consumer.accept(builder);
+		ExecutionInput executionInput = ExecutionInput.newExecutionInput("{}").build();
+		ExecutionResult result = builder.build();
+
+		when(this.transport.execute(this.requestCaptor.capture()))
+				.thenReturn(Mono.just(new RequestOutput(executionInput, result)));
+	}
+
+	private void serialize(String data, ExecutionResultImpl.Builder builder) {
+		try {
+			builder.data(OBJECT_MAPPER.readValue(data, Map.class));
+		}
+		catch (JsonProcessingException ex) {
+			throw new IllegalStateException(ex);
+		}
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/client/GraphQlClientTests.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright 2002-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.graphql.client;
+
+import java.util.List;
+import java.util.Map;
+
+import graphql.GraphQLError;
+import graphql.GraphqlErrorBuilder;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.core.ParameterizedTypeReference;
+import org.springframework.graphql.GraphQlRequest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+
+/**
+ * Test {@link GraphQlClient} with a mock/test transport.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class GraphQlClientTests extends GraphQlClientTestSupport {
+
+	@Test
+	void entity() {
+		String document = "{me {name}}";
+		setMockResponse("{\"me\": {\"name\":\"Luke Skywalker\"}}");
+
+		GraphQlClient.ResponseSpec spec = execute(document);
+
+		MovieCharacter luke = MovieCharacter.create("Luke Skywalker");
+		assertThat(spec.toEntity("me", MovieCharacter.class)).isEqualTo(luke);
+
+		Map<String, MovieCharacter> map = spec.toEntity("", new ParameterizedTypeReference<Map<String, MovieCharacter>>() {});
+		assertThat(map).containsEntry("me", luke);
+
+		assertThat(request().getDocument()).contains(document);
+	}
+
+	@Test
+	void entityList() {
+
+		String document = "{me {name, friends}}";
+		setMockResponse("{" +
+				"  \"me\":{" +
+				"      \"name\":\"Luke Skywalker\","
+				+ "      \"friends\":[{\"name\":\"Han Solo\"}, {\"name\":\"Leia Organa\"}]" +
+				"  }" +
+				"}");
+
+		GraphQlClient.ResponseSpec spec = execute(document);
+
+		MovieCharacter han = MovieCharacter.create("Han Solo");
+		MovieCharacter leia = MovieCharacter.create("Leia Organa");
+
+		List<MovieCharacter> characters = spec.toEntityList("me.friends", MovieCharacter.class);
+		assertThat(characters).containsExactly(han, leia);
+
+		characters = spec.toEntityList("me.friends", new ParameterizedTypeReference<MovieCharacter>() {});
+		assertThat(characters).containsExactly(han, leia);
+
+		assertThat(request().getDocument()).contains(document);
+	}
+
+	@Test
+	void operationNameAndVariables() {
+
+		String document = "query HeroNameAndFriends($episode: Episode) {" +
+				"  hero(episode: $episode) {" +
+				"    name"
+				+ "  }" +
+				"}";
+		setMockResponse("{\"hero\": {\"name\":\"R2-D2\"}}");
+
+		GraphQlClient.ResponseSpec spec = graphQlClient().document(document)
+				.operationName("HeroNameAndFriends")
+				.variable("episode", "JEDI")
+				.variable("foo", "bar")
+				.variable("keyOnly", null)
+				.execute()
+				.block(TIMEOUT);
+
+		assertThat(spec).isNotNull();
+
+		MovieCharacter character = spec.toEntity("hero", MovieCharacter.class);
+		assertThat(character).isEqualTo(MovieCharacter.create("R2-D2"));
+
+		GraphQlRequest request = request();
+		assertThat(request.getDocument()).contains(document);
+		assertThat(request.getOperationName()).isEqualTo("HeroNameAndFriends");
+		assertThat(request.getVariables()).hasSize(3);
+		assertThat(request.getVariables()).containsEntry("episode", "JEDI");
+		assertThat(request.getVariables()).containsEntry("foo", "bar");
+		assertThat(request.getVariables()).containsEntry("keyOnly", null);
+	}
+
+	@Test
+	void errors() {
+
+		String document = "{me {name, friends}}";
+		setMockResponse(
+				GraphqlErrorBuilder.newError().message("some error").build(),
+				GraphqlErrorBuilder.newError().message("some other error").build());
+
+		GraphQlClient.ResponseSpec spec = execute(document);
+
+		assertThat(spec.errors()).extracting(GraphQLError::getMessage)
+				.containsExactly("some error", "some other error");
+	}
+
+	private GraphQlClient.ResponseSpec execute(String document) {
+		GraphQlClient.ResponseSpec spec = graphQlClient().document(document).execute().block(TIMEOUT);
+		assertThat(spec).isNotNull();
+		return spec;
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/client/MockGraphQlWebSocketServer.java

@@ -22,6 +22,8 @@ import java.util.function.Function;
 
 import graphql.ExecutionResult;
 import graphql.GraphQLError;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -33,13 +35,16 @@ import org.springframework.web.reactive.socket.WebSocketHandler;
 import org.springframework.web.reactive.socket.WebSocketSession;
 
 /**
- * GraphQL over WebSocket handler to use as a server-side
+ * GraphQL over WebSocket {@link WebSocketHandler} to use as a server-side
  * {@link WebSocketHandler} that is configured with expected requests and
  * the responses to send.
  *
  * @author Rossen Stoyanchev
  */
-public class MockWebSocketServer implements WebSocketHandler {
+public final class MockGraphQlWebSocketServer implements WebSocketHandler {
+
+	private final static Log logger = LogFactory.getLog(MockGraphQlWebSocketServer.class);
+
 
 	@Nullable
 	private Function<Map<String, Object>, Mono<Object>> connectionInitHandler;
@@ -73,7 +78,8 @@ public class MockWebSocketServer implements WebSocketHandler {
 		return session.send(session.receive()
 				.map(codecDelegate::decode)
 				.flatMap(this::handleMessage)
-				.map(message -> codecDelegate.encode(session, message)));
+				.map(message -> codecDelegate.encode(session, message)))
+				.doOnError(ex -> logger.error("Session handling error: " + ex.getMessage(), ex));
 	}
 
 	@SuppressWarnings("SuspiciousMethodCalls")

spring-graphql/src/test/java/org/springframework/graphql/client/MockWebSocketGraphQlTransportTests.java

@@ -37,6 +37,7 @@ import org.springframework.graphql.GraphQlRequest;
 import org.springframework.graphql.support.MapExecutionResult;
 import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage;
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.codec.ClientCodecConfigurer;
 import org.springframework.web.reactive.socket.CloseStatus;
 import org.springframework.web.reactive.socket.WebSocketHandler;
 import org.springframework.web.reactive.socket.WebSocketSession;
@@ -48,7 +49,9 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 /**
- * Unit tests for {@link WebSocketGraphQlTransport}.
+ * Tests for {@link WebSocketGraphQlTransport} using {@link TestWebSocketClient}
+ * and {@link MockGraphQlWebSocketServer}.
+ *
  * @author Rossen Stoyanchev
  */
 public class MockWebSocketGraphQlTransportTests {
@@ -58,12 +61,11 @@ public class MockWebSocketGraphQlTransportTests {
 	private static final WebSocketCodecDelegate CODEC_DELEGATE = new WebSocketCodecDelegate();
 
 
-	private final MockWebSocketServer mockServer = new MockWebSocketServer();
+	private final MockGraphQlWebSocketServer mockServer = new MockGraphQlWebSocketServer();
 
-	private final TestWebSocketClient testClient = new TestWebSocketClient(this.mockServer);
+	private final TestWebSocketClient webSocketClient = new TestWebSocketClient(this.mockServer);
 	
-	private final WebSocketGraphQlTransport transport = 
-			WebSocketGraphQlTransport.builder(URI.create("/"), this.testClient).build();
+	private final WebSocketGraphQlTransport transport = createTransport(this.webSocketClient);
 
 	private final ExecutionResult result1 = MapExecutionResult.forDataOnly(Collections.singletonMap("key1", "value1"));
 
@@ -72,7 +74,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void request() {
-		GraphQlRequest request = this.mockServer.expectOperation("Query1").andRespond(this.result1);
+		GraphQlRequest request = this.mockServer.expectOperation("{Query1}").andRespond(this.result1);
 
 		StepVerifier.create(this.transport.execute(request))
 				.expectNext(this.result1).expectComplete()
@@ -85,7 +87,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void requestStream() {
-		GraphQlRequest request = this.mockServer.expectOperation("Sub1").andStream(Flux.just(this.result1, result2));
+		GraphQlRequest request = this.mockServer.expectOperation("{Sub1}").andStream(Flux.just(this.result1, result2));
 
 		StepVerifier.create(this.transport.executeSubscription(request))
 				.expectNext(this.result1, result2).expectComplete()
@@ -98,7 +100,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void requestError() {
-		GraphQlRequest request = this.mockServer.expectOperation("Query1")
+		GraphQlRequest request = this.mockServer.expectOperation("{Query1}")
 				.andRespondWithError(GraphqlErrorBuilder.newError().message("boo").build());
 
 		StepVerifier.create(this.transport.execute(request))
@@ -116,7 +118,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void requestStreamError() {
-		GraphQlRequest request = this.mockServer.expectOperation("Sub1")
+		GraphQlRequest request = this.mockServer.expectOperation("{Sub1}")
 				.andStreamWithError(Flux.just(this.result1), GraphqlErrorBuilder.newError().message("boo").build());
 
 		StepVerifier.create(this.transport.executeSubscription(request))
@@ -134,7 +136,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void requestCancelled() {
-		GraphQlRequest request = this.mockServer.expectOperation("Query1").andRespond(Mono.never());
+		GraphQlRequest request = this.mockServer.expectOperation("{Query1}").andRespond(Mono.never());
 
 		StepVerifier.create(this.transport.execute(request))
 				.thenAwait(Duration.ofMillis(200))
@@ -148,7 +150,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void requestStreamCancelled() {
-		GraphQlRequest request = this.mockServer.expectOperation("s1")
+		GraphQlRequest request = this.mockServer.expectOperation("{Sub1}")
 				.andStream(Flux.just(this.result1).concatWith(Flux.never()));
 
 		StepVerifier.create(this.transport.executeSubscription(request))
@@ -165,17 +167,16 @@ public class MockWebSocketGraphQlTransportTests {
 
 	@Test
 	void start() {
-		MockWebSocketServer handler = new MockWebSocketServer();
+		MockGraphQlWebSocketServer handler = new MockGraphQlWebSocketServer();
 		handler.connectionInitHandler(payload -> Mono.just(Collections.singletonMap("key", payload.get("key") + "Ack")));
 
 		TestWebSocketClient client = new TestWebSocketClient(handler);
 		Map<String, String> initPayload = Collections.singletonMap("key", "valueInit");
 		AtomicReference<Map<String, Object>> connectionAckRef = new AtomicReference<>();
 
-		WebSocketGraphQlTransport transport = WebSocketGraphQlTransport.builder(URI.create("/"), client)
-				.connectionInitPayload(initPayload)
-				.connectionAckHandler(connectionAckRef::set)
-				.build();
+		WebSocketGraphQlTransport transport = new WebSocketGraphQlTransport(
+				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(),
+				initPayload, connectionAckRef::set);
 
 		transport.start().block(TIMEOUT);
 
@@ -189,27 +190,27 @@ public class MockWebSocketGraphQlTransportTests {
 
 		// Start
 		this.transport.start().block(TIMEOUT);
-		assertThat(this.testClient.getConnectionCount()).isEqualTo(1);
-		assertThat(this.testClient.getConnection(0).isOpen()).isTrue();
+		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(1);
+		assertThat(this.webSocketClient.getConnection(0).isOpen()).isTrue();
 
 		// Stop
 		this.transport.stop().block(TIMEOUT);
-		assertThat(this.testClient.getConnection(0).isOpen()).isFalse();
-		assertThat(this.testClient.getConnection(0).closeStatus().block(TIMEOUT)).isEqualTo(CloseStatus.NORMAL);
+		assertThat(this.webSocketClient.getConnection(0).isOpen()).isFalse();
+		assertThat(this.webSocketClient.getConnection(0).closeStatus().block(TIMEOUT)).isEqualTo(CloseStatus.NORMAL);
 
 		// New requests are rejected
-		GraphQlRequest request = this.mockServer.expectOperation("Query1").andRespond(this.result1);
+		GraphQlRequest request = this.mockServer.expectOperation("{Query1}").andRespond(this.result1);
 		StepVerifier.create(this.transport.execute(request))
 				.expectErrorMessage("WebSocketGraphQlTransport has been stopped")
 				.verify(TIMEOUT);
 
 		// Start
 		this.transport.start().block(TIMEOUT);
-		assertThat(this.testClient.getConnectionCount()).isEqualTo(2);
-		assertThat(this.testClient.getConnection(1).isOpen()).isTrue();
+		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(2);
+		assertThat(this.webSocketClient.getConnection(1).isOpen()).isTrue();
 
 		// Requests allowed again
-		request = this.mockServer.expectOperation("Query1").andRespond(this.result1);
+		request = this.mockServer.expectOperation("{Query1}").andRespond(this.result1);
 		StepVerifier.create(this.transport.execute(request))
 				.expectNext(this.result1).expectComplete()
 				.verify(TIMEOUT);
@@ -218,26 +219,26 @@ public class MockWebSocketGraphQlTransportTests {
 	@Test
 	void sessionIsCachedUntilClosed() {
 
-		GraphQlRequest request1 = this.mockServer.expectOperation("Query1").andRespond(this.result1);
+		GraphQlRequest request1 = this.mockServer.expectOperation("{Query1}").andRespond(this.result1);
 		StepVerifier.create(this.transport.execute(request1)).expectNext(this.result1).expectComplete().verify(TIMEOUT);
 
-		assertThat(this.testClient.getConnectionCount()).isEqualTo(1);
-		TestWebSocketConnection originalConnection = this.testClient.getConnection(0);
+		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(1);
+		TestWebSocketConnection originalConnection = this.webSocketClient.getConnection(0);
 
-		GraphQlRequest request2 = this.mockServer.expectOperation("Query2").andRespond(this.result2);
+		GraphQlRequest request2 = this.mockServer.expectOperation("{Query2}").andRespond(this.result2);
 		StepVerifier.create(this.transport.execute(request2)).expectNext(this.result2).expectComplete().verify(TIMEOUT);
 
-		assertThat(this.testClient.getConnectionCount()).isEqualTo(1);
-		assertThat(this.testClient.getConnection(0)).isSameAs(originalConnection);
+		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(1);
+		assertThat(this.webSocketClient.getConnection(0)).isSameAs(originalConnection);
 
 		// Close the connection
 		originalConnection.closeServerSession(CloseStatus.NORMAL).block(TIMEOUT);
 
-		request1 = this.mockServer.expectOperation("Query1").andRespond(this.result1);
+		request1 = this.mockServer.expectOperation("{Query1}").andRespond(this.result1);
 		StepVerifier.create(this.transport.execute(request1)).expectNext(this.result1).expectComplete().verify(TIMEOUT);
 
-		assertThat(this.testClient.getConnectionCount()).isEqualTo(2);
-		assertThat(this.testClient.getConnection(1)).isNotSameAs(originalConnection);
+		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(2);
+		assertThat(this.webSocketClient.getConnection(1)).isNotSameAs(originalConnection);
 	}
 
 	@Test
@@ -260,7 +261,7 @@ public class MockWebSocketGraphQlTransportTests {
 
 		// Errors before GraphQL session initialized should be routed, no hanging on start
 
-		MockWebSocketServer handler = new MockWebSocketServer();
+		MockGraphQlWebSocketServer handler = new MockGraphQlWebSocketServer();
 		handler.connectionInitHandler(initPayload -> Mono.error(new IllegalStateException("boo")));
 
 		TestWebSocketClient client = new TestWebSocketClient(handler);
@@ -278,20 +279,20 @@ public class MockWebSocketGraphQlTransportTests {
 		TestWebSocketClient client = new TestWebSocketClient(new UnexpectedResponseHandler());
 		WebSocketGraphQlTransport transport = createTransport(client);
 
-		String expectedMessage = "GraphQlSession over client-session-1 disconnected " +
-				"with CloseStatus[code=1002, reason=null]";
+		String expectedMessage = "disconnected with CloseStatus[code=1002, reason=null]";
 
-		StepVerifier.create(transport.execute(new GraphQlRequest("Query1")))
-				.expectErrorMessage(expectedMessage)
+		StepVerifier.create(transport.execute(new GraphQlRequest("{Query1}")))
+				.expectErrorSatisfies(ex -> assertThat(ex).hasMessageEndingWith(expectedMessage))
 				.verify(TIMEOUT);
 	}
 
-	private WebSocketGraphQlTransport createTransport(WebSocketClient client) {
-		return WebSocketGraphQlTransport.builder(URI.create("/"), client).build();
+	private static WebSocketGraphQlTransport createTransport(WebSocketClient client) {
+		return new WebSocketGraphQlTransport(
+				URI.create("/"), HttpHeaders.EMPTY, client, ClientCodecConfigurer.create(), null, p -> {});
 	}
 
 	private void assertActualClientMessages(GraphQlWebSocketMessage... expectedMessages) {
-		assertActualClientMessages(this.testClient.getConnection(0), expectedMessages);
+		assertActualClientMessages(this.webSocketClient.getConnection(0), expectedMessages);
 	}
 
 	private void assertActualClientMessages(

spring-graphql/src/test/java/org/springframework/graphql/client/MovieCharacter.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2002-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.client;
+
+import org.springframework.lang.Nullable;
+import org.springframework.util.ObjectUtils;
+
+
+public class MovieCharacter {
+
+	@Nullable
+	private String name;
+
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	@Nullable
+	public String getName() {
+		return this.name;
+	}
+
+
+	public static MovieCharacter create(String name) {
+		MovieCharacter character = new MovieCharacter();
+		character.setName(name);
+		return character;
+	}
+
+
+	@Override
+	public boolean equals(Object other) {
+		if (this == other) {
+			return true;
+		}
+		if (other == null || getClass() != other.getClass()) {
+			return false;
+		}
+		return ObjectUtils.nullSafeEquals(this.name, ((MovieCharacter) other).name);
+	}
+
+	@Override
+	public int hashCode() {
+		return (this.name != null) ? this.name.hashCode() : super.hashCode();
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/client/TestWebSocketClient.java

@@ -37,7 +37,7 @@ import org.springframework.web.reactive.socket.client.WebSocketClient;
  *
  * @author Rossen Stoyanchev
  */
-public class TestWebSocketClient implements WebSocketClient {
+final class TestWebSocketClient implements WebSocketClient {
 
 	private final WebSocketHandler serverHandler;
 

spring-graphql/src/test/java/org/springframework/graphql/client/TestWebSocketConnection.java

@@ -49,11 +49,15 @@ import org.springframework.web.reactive.socket.adapter.AbstractWebSocketSession;
  *
  * @author Rossen Stoyanchev
  */
-public class TestWebSocketConnection {
+final class TestWebSocketConnection {
 
 	private static final AtomicLong connectionIndex = new AtomicLong();
 
 
+	private final URI url;
+
+	private final HttpHeaders headers;
+
 	private final TestWebSocketSession clientSession;
 
 	private final TestWebSocketSession serverSession;
@@ -61,6 +65,9 @@ public class TestWebSocketConnection {
 
 	public TestWebSocketConnection(URI url, HttpHeaders headers) {
 
+		this.url = url;
+		this.headers = headers;
+
 		long id = connectionIndex.incrementAndGet();
 
 		Sinks.Many<WebSocketMessage> clientSink = Sinks.many().unicast().onBackpressureBuffer();
@@ -77,6 +84,14 @@ public class TestWebSocketConnection {
 	}
 
 
+	public URI getUrl() {
+		return this.url;
+	}
+
+	public HttpHeaders getHeaders() {
+		return this.headers;
+	}
+
 	/**
 	 * Return {@code true} if both client and server sessions are open.
 	 */

spring-graphql/src/test/java/org/springframework/graphql/client/WebGraphQlClientBuilderTests.java

@@ -0,0 +1,230 @@
+/*
+ * Copyright 2002-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.client;
+
+import java.net.URI;
+import java.time.Duration;
+import java.util.stream.Stream;
+
+import graphql.ExecutionInput;
+import graphql.ExecutionResultImpl;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import reactor.core.publisher.Mono;
+
+import org.springframework.graphql.RequestOutput;
+import org.springframework.graphql.support.DocumentSource;
+import org.springframework.graphql.web.WebGraphQlHandler;
+import org.springframework.graphql.web.WebInput;
+import org.springframework.graphql.web.WebInterceptor;
+import org.springframework.graphql.web.WebOutput;
+import org.springframework.graphql.web.webflux.GraphQlHttpHandler;
+import org.springframework.graphql.web.webflux.GraphQlWebSocketHandler;
+import org.springframework.http.codec.ClientCodecConfigurer;
+import org.springframework.http.server.reactive.HttpHandler;
+import org.springframework.test.web.reactive.server.HttpHandlerConnector;
+import org.springframework.web.reactive.function.client.WebClient;
+import org.springframework.web.reactive.function.server.HandlerStrategies;
+import org.springframework.web.reactive.function.server.RouterFunction;
+import org.springframework.web.reactive.function.server.RouterFunctions;
+import org.springframework.web.reactive.function.server.ServerResponse;
+import org.springframework.web.reactive.socket.WebSocketHandler;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.springframework.web.reactive.function.server.RouterFunctions.route;
+
+/**
+ * Tests for the builders of Web {@code GraphQlClient} extensions, using a
+ * {@link WebInterceptor} to capture the WebInput on the server side, and return
+ * with no handling.
+ *
+ * <ul>
+ * <li>{@link HttpGraphQlClient} via {@link HttpHandlerConnector} to {@link GraphQlHttpHandler}
+ * <li>{@link WebSocketGraphQlClient} via a {@link TestWebSocketConnection} to {@link GraphQlWebSocketHandler}
+ * </ul>
+ *
+ * @author Rossen Stoyanchev
+ */
+public class WebGraphQlClientBuilderTests {
+
+	private static final String DOCUMENT = "{ Query }";
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+
+	public static Stream<ClientBuilderSetup> argumentSource() {
+		return Stream.of(new HttpBuilderSetup(), new WebSocketBuilderSetup());
+	}
+
+
+	@ParameterizedTest
+	@MethodSource("argumentSource")
+	void mutateUrlHeaders(ClientBuilderSetup builderSetup) {
+
+		String url = "/graphql-one";
+
+		// Original
+		WebGraphQlClient.Builder<?> builder = builderSetup.initBuilder()
+				.url(url)
+				.headers(headers -> headers.add("h", "one"));
+
+		WebGraphQlClient client = builder.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+
+		WebInput input = builderSetup.getWebInput();
+		assertThat(input.getUri().toString()).isEqualTo(url);
+		assertThat(input.getHeaders().get("h")).containsExactly("one");
+
+		// Mutate to add header value
+		builder = client.mutate().headers(headers -> headers.add("h", "two"));
+		client = builder.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+		assertThat(builderSetup.getWebInput().getHeaders().get("h")).containsExactly("one", "two");
+
+		// Mutate to replace header
+		builder = client.mutate().header("h", "three", "four");
+		client = builder.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+
+		input = builderSetup.getWebInput();
+		assertThat(input.getUri().toString()).isEqualTo(url);
+		assertThat(input.getHeaders().get("h")).containsExactly("three", "four");
+	}
+
+	@Test
+	void mutateWebTestClientViaConsumer() {
+		HttpBuilderSetup clientSetup = new HttpBuilderSetup();
+
+		// Original header value
+		HttpGraphQlClient.Builder<?> builder = clientSetup.initBuilder()
+				.webClient(testClientBuilder -> testClientBuilder.defaultHeaders(h -> h.add("h", "one")));
+
+		HttpGraphQlClient client = builder.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+		assertThat(clientSetup.getWebInput().getHeaders().get("h")).containsExactly("one");
+
+		// Mutate to add header value
+		HttpGraphQlClient.Builder<?> builder2 = client.mutate()
+				.webClient(testClientBuilder -> testClientBuilder.defaultHeaders(h -> h.add("h", "two")));
+
+		client = builder2.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+		assertThat(clientSetup.getWebInput().getHeaders().get("h")).containsExactly("one", "two");
+
+		// Mutate to replace header
+		HttpGraphQlClient.Builder<?> builder3 = client.mutate()
+				.webClient(testClientBuilder -> testClientBuilder.defaultHeader("h", "three"));
+
+		client = builder3.build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+		assertThat(clientSetup.getWebInput().getHeaders().get("h")).containsExactly("three");
+	}
+
+	@ParameterizedTest
+	@MethodSource("argumentSource")
+	void mutateDocumentSource(ClientBuilderSetup builderSetup) {
+
+		DocumentSource documentSource = name -> name.equals("name") ?
+				Mono.just(DOCUMENT) : Mono.error(new IllegalArgumentException());
+
+		// Original
+		WebGraphQlClient.Builder<?> builder = builderSetup.initBuilder().documentSource(documentSource);
+		WebGraphQlClient client = builder.build();
+		client.documentName("name").execute().block(TIMEOUT);
+
+		WebInput input = builderSetup.getWebInput();
+		assertThat(input.getDocument()).isEqualTo(DOCUMENT);
+
+		// Mutate
+		client = client.mutate().build();
+		client.documentName("name").execute().block(TIMEOUT);
+
+		input = builderSetup.getWebInput();
+		assertThat(input.getDocument()).isEqualTo(DOCUMENT);
+	}
+
+	@ParameterizedTest
+	@MethodSource("argumentSource")
+	void url(ClientBuilderSetup builderSetup) {
+
+		WebGraphQlClient client = builderSetup.initBuilder().url("/graphql one").build();
+		client.document(DOCUMENT).execute().block(TIMEOUT);
+
+		assertThat(builderSetup.getWebInput().getUri().toString()).isEqualTo("/graphql%20one");
+	}
+
+
+
+	private interface ClientBuilderSetup {
+
+		WebGraphQlClient.Builder<?> initBuilder();
+
+		WebInput getWebInput();
+
+	}
+
+
+	private abstract static class AbstractBuilderSetup implements ClientBuilderSetup {
+
+		private WebInput webInput;
+
+		protected WebGraphQlHandler webGraphQlHandler() {
+			return WebGraphQlHandler.builder(requestInput -> Mono.error(new UnsupportedOperationException()))
+					.interceptor((input, chain) -> {
+						this.webInput = input;
+						return Mono.just(new WebOutput(new RequestOutput(
+								ExecutionInput.newExecutionInput().query("{ notUsed }").build(),
+								ExecutionResultImpl.newExecutionResult().build())));
+					})
+					.build();
+		}
+
+		@Override
+		public WebInput getWebInput() {
+			return this.webInput;
+		}
+
+	}
+
+
+	private static class HttpBuilderSetup extends AbstractBuilderSetup {
+
+		@Override
+		public HttpGraphQlClient.Builder<?> initBuilder() {
+			GraphQlHttpHandler handler = new GraphQlHttpHandler(webGraphQlHandler());
+			RouterFunction<ServerResponse> routerFunction = route().POST("/**", handler::handleRequest).build();
+			HttpHandler httpHandler = RouterFunctions.toHttpHandler(routerFunction, HandlerStrategies.withDefaults());
+			HttpHandlerConnector connector = new HttpHandlerConnector(httpHandler);
+			return HttpGraphQlClient.builder(WebClient.builder().clientConnector(connector));
+		}
+
+	}
+
+
+	private static class WebSocketBuilderSetup extends AbstractBuilderSetup {
+
+		@Override
+		public WebSocketGraphQlClient.Builder<?> initBuilder() {
+			ClientCodecConfigurer configurer = ClientCodecConfigurer.create();
+			WebSocketHandler handler = new GraphQlWebSocketHandler(webGraphQlHandler(), configurer, Duration.ofSeconds(5));
+			return WebSocketGraphQlClient.builder(URI.create(""), new TestWebSocketClient(handler));
+		}
+
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/client/WebSocketGraphQlClientBuilderTests.java

@@ -1,103 +0,0 @@
-/*
- * Copyright 2002-2022 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.
- * You may obtain a copy of the License at
- *
- *      https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.springframework.graphql.client;
-
-import java.net.URI;
-import java.time.Duration;
-import java.util.Collections;
-
-import graphql.ExecutionResult;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-
-import org.springframework.graphql.support.MapExecutionResult;
-
-import static org.assertj.core.api.Assertions.assertThat;
-
-/**
- * Unit tests for {@link GraphQlClient.Builder}.
- * @author Rossen Stoyanchev
- */
-public class WebSocketGraphQlClientBuilderTests {
-
-	private TestWebSocketClient webSocketClient;
-
-
-	@BeforeEach
-	void setUp() {
-		MockGraphQlWebSocketServer mockServer = new MockGraphQlWebSocketServer();
-		this.webSocketClient = new TestWebSocketClient(mockServer);
-
-		ExecutionResult result = MapExecutionResult.forDataOnly(Collections.singletonMap("key1", "value1"));
-		mockServer.expectOperation("{Query1}").andRespond(result);
-	}
-
-
-	@Test
-	void mutate() {
-
-		// Original
-
-		URI url = URI.create("/graphql");
-
-		WebSocketGraphQlClient graphQlClient = WebSocketGraphQlClient.builder(this.webSocketClient)
-				.url(url)
-				.header("header1", "value1")
-				.header("header2", "value2")
-				.build();
-
-		graphQlClient.document("{Query1}").execute().block(Duration.ofSeconds(5));
-
-		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(1);
-		assertThat(this.webSocketClient.getConnection(0).getUrl()).isEqualTo(url);
-		assertThat(this.webSocketClient.getConnection(0).getHeaders()).hasSize(2)
-				.containsEntry("header1", Collections.singletonList("value1"))
-				.containsEntry("header2", Collections.singletonList("value2"));
-
-		// Mutate
-
-		URI anotherUrl = URI.create("/another-graphql");
-
-		WebSocketGraphQlClient anotherClient = graphQlClient.mutate()
-				.url(anotherUrl)
-				.headers(headers -> {
-					headers.set("header1", "anotherValue1");
-					headers.set("header2", "anotherValue2");
-				})
-				.build();
-
-		anotherClient.document("{Query1}").execute().block(Duration.ofSeconds(5));
-
-		assertThat(this.webSocketClient.getConnectionCount()).isEqualTo(2);
-		assertThat(this.webSocketClient.getConnection(1).getUrl()).isEqualTo(anotherUrl);
-		assertThat(this.webSocketClient.getConnection(1).getHeaders()).hasSize(2)
-				.containsEntry("header1", Collections.singletonList("anotherValue1"))
-				.containsEntry("header2", Collections.singletonList("anotherValue2"));
-
-		// Original not affected (stop + start original client, to connect again)
-
-		graphQlClient.stop().block(Duration.ofSeconds(5));
-		graphQlClient.start().block(Duration.ofSeconds(5));
-		graphQlClient.document("{Query1}").execute().block();
-
-		assertThat(this.webSocketClient.getConnection(0).getUrl()).isEqualTo(url);
-		assertThat(this.webSocketClient.getConnection(0).getHeaders()).hasSize(2)
-				.containsEntry("header1", Collections.singletonList("value1"))
-				.containsEntry("header2", Collections.singletonList("value2"));
-	}
-
-}

spring-graphql/src/test/java/org/springframework/graphql/client/WebSocketGraphQlClientTests.java

@@ -1,135 +0,0 @@
-/*
- * Copyright 2002-2022 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.
- * You may obtain a copy of the License at
- *
- *      https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.springframework.graphql.client;
-
-import java.net.URI;
-import java.time.Duration;
-import java.util.List;
-
-import org.junit.jupiter.api.Test;
-import reactor.core.publisher.Flux;
-import reactor.test.StepVerifier;
-
-import org.springframework.context.annotation.AnnotationConfigApplicationContext;
-import org.springframework.graphql.Book;
-import org.springframework.graphql.BookCriteria;
-import org.springframework.graphql.BookSource;
-import org.springframework.graphql.GraphQlSetup;
-import org.springframework.graphql.data.method.annotation.Argument;
-import org.springframework.graphql.data.method.annotation.QueryMapping;
-import org.springframework.graphql.data.method.annotation.SubscriptionMapping;
-import org.springframework.graphql.web.WebGraphQlHandler;
-import org.springframework.graphql.web.webflux.GraphQlWebSocketHandler;
-import org.springframework.http.codec.ClientCodecConfigurer;
-import org.springframework.stereotype.Controller;
-import org.springframework.web.reactive.socket.WebSocketHandler;
-
-import static org.assertj.core.api.Assertions.assertThat;
-
-/**
- * Test making requests through {@link WebSocketGraphQlClient} to
- * {@link GraphQlWebSocketHandler} via a mock WebSocket connection.
- *
- * @author Rossen Stoyanchev
- */
-public class WebSocketGraphQlClientTests {
-
-	private final TestWebSocketClient webSocketClient = initWebSocketClient();
-
-
-	@Test
-	void query() {
-		String document = "{ " +
-				"  booksByCriteria(criteria: {author:\"Orwell\"}) { " +
-				"    id" +
-				"    name" +
-				"  }" +
-				"}";
-
-		List<Book> books = WebSocketGraphQlClient.create(URI.create("/"), this.webSocketClient)
-				.document(document)
-				.execute()
-				.map(response -> response.toEntityList("booksByCriteria", Book.class))
-				.block(Duration.ofSeconds(5));
-
-		assertThat(books).hasSize(2);
-		assertThat(books.get(0).getName()).isEqualTo("Nineteen Eighty-Four");
-		assertThat(books.get(1).getName()).isEqualTo("Animal Farm");
-	}
-
-	@Test
-	void subscription() {
-		String document = "subscription { " +
-				"  bookSearch(author:\"Orwell\") { " +
-				"    id" +
-				"    name" +
-				"  }" +
-				"}";
-
-		Flux<Book> bookFlux = WebSocketGraphQlClient.create(URI.create("/"), this.webSocketClient)
-				.document(document)
-				.executeSubscription()
-				.map(response -> response.toEntity("bookSearch", Book.class));
-
-		StepVerifier.create(bookFlux)
-				.consumeNextWith(book -> {
-					assertThat(book.getId()).isEqualTo(1);
-					assertThat(book.getName()).isEqualTo("Nineteen Eighty-Four");
-				})
-				.consumeNextWith(book -> {
-					assertThat(book.getId()).isEqualTo(5);
-					assertThat(book.getName()).isEqualTo("Animal Farm");
-				})
-				.verifyComplete();
-	}
-
-
-	private TestWebSocketClient initWebSocketClient() {
-
-		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
-		context.register(BookController.class);
-		context.refresh();
-
-		WebGraphQlHandler webGraphQlHandler = GraphQlSetup.schemaResource(BookSource.schema)
-				.runtimeWiringForAnnotatedControllers(context)
-				.toWebGraphQlHandler();
-
-		WebSocketHandler webSocketServerHandler = new GraphQlWebSocketHandler(
-				webGraphQlHandler, ClientCodecConfigurer.create(), Duration.ofSeconds(5));
-
-		return new TestWebSocketClient(webSocketServerHandler);
-	}
-
-
-
-	@SuppressWarnings("unused")
-	@Controller
-	private static class BookController {
-
-		@QueryMapping
-		public List<Book> booksByCriteria(@Argument BookCriteria criteria) {
-			return BookSource.findBooksByAuthor(criteria.getAuthor());
-		}
-
-		@SubscriptionMapping
-		public Flux<Book> bookSearch(@Argument String author) {
-			return Flux.fromIterable(BookSource.findBooksByAuthor(author)).delayElements(Duration.ofMillis(50));
-		}
-
-	}
-
-}