Extract GraphQlWebSocketMessage and WebSocketCodecDelegate

Allows some reuse between WebFlux and WebMVC and between client and server.

See gh-10

spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java

@@ -19,14 +19,13 @@ package org.springframework.graphql.web.webflux;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-import graphql.ErrorType;
 import graphql.ExecutionResult;
+import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -35,23 +34,12 @@ import org.reactivestreams.Subscription;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import org.springframework.core.ParameterizedTypeReference;
-import org.springframework.core.ResolvableType;
-import org.springframework.core.codec.Decoder;
-import org.springframework.core.codec.Encoder;
-import org.springframework.core.io.buffer.DataBuffer;
-import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.graphql.web.WebGraphQlHandler;
 import org.springframework.graphql.web.WebInput;
 import org.springframework.graphql.web.WebOutput;
-import org.springframework.http.MediaType;
-import org.springframework.http.codec.DecoderHttpMessageReader;
-import org.springframework.http.codec.EncoderHttpMessageWriter;
-import org.springframework.http.codec.ServerCodecConfigurer;
-import org.springframework.lang.Nullable;
+import org.springframework.http.codec.CodecConfigurer;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
-import org.springframework.util.MimeTypeUtils;
 import org.springframework.web.reactive.socket.CloseStatus;
 import org.springframework.web.reactive.socket.HandshakeInfo;
 import org.springframework.web.reactive.socket.WebSocketHandler;
@@ -72,57 +60,36 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 
 	private static final List<String> SUB_PROTOCOL_LIST = Arrays.asList("graphql-transport-ws", "graphql-ws");
 
-	static final ResolvableType MAP_RESOLVABLE_TYPE =
-			ResolvableType.forType(new ParameterizedTypeReference<Map<String, Object>>() {});
-
 
 	private final WebGraphQlHandler graphQlHandler;
 
-	private final Decoder<?> decoder;
-
-	private final Encoder<?> encoder;
+	private final WebSocketCodecDelegate codecDelegate;
 
 	private final Duration initTimeoutDuration;
 
+
 	/**
 	 * Create a new instance.
 	 * @param graphQlHandler common handler for GraphQL over WebSocket requests
-	 * @param configurer codec configurer for JSON encoding and decoding
+	 * @param codecConfigurer codec configurer for JSON encoding and decoding
 	 * @param connectionInitTimeout the time within which the {@code CONNECTION_INIT} type
 	 * message must be received.
 	 */
 	public GraphQlWebSocketHandler(
-			WebGraphQlHandler graphQlHandler, ServerCodecConfigurer configurer,
-			Duration connectionInitTimeout) {
+			WebGraphQlHandler graphQlHandler, CodecConfigurer codecConfigurer, Duration connectionInitTimeout) {
 
 		Assert.notNull(graphQlHandler, "WebGraphQlHandler is required");
 		this.graphQlHandler = graphQlHandler;
-		this.decoder = initDecoder(configurer);
-		this.encoder = initEncoder(configurer);
+		this.codecDelegate = new WebSocketCodecDelegate(codecConfigurer);
 		this.initTimeoutDuration = connectionInitTimeout;
 	}
 
-	private static Decoder<?> initDecoder(ServerCodecConfigurer configurer) {
-		return configurer.getReaders().stream()
-				.filter((reader) -> reader.canRead(MAP_RESOLVABLE_TYPE, MediaType.APPLICATION_JSON))
-				.map((reader) -> ((DecoderHttpMessageReader<?>) reader).getDecoder())
-				.findFirst()
-				.orElseThrow(() -> new IllegalArgumentException("No JSON Decoder"));
-	}
 
-	private static Encoder<?> initEncoder(ServerCodecConfigurer configurer) {
-		return configurer.getWriters().stream()
-				.filter((writer) -> writer.canWrite(MAP_RESOLVABLE_TYPE, MediaType.APPLICATION_JSON))
-				.map((writer) -> ((EncoderHttpMessageWriter<?>) writer).getEncoder())
-				.findFirst()
-				.orElseThrow(() -> new IllegalArgumentException("No JSON Encoder"));
-	}
-
-	@Override
 	public List<String> getSubProtocols() {
 		return SUB_PROTOCOL_LIST;
 	}
 
+
 	@Override
 	public Mono<Void> handle(WebSocketSession session) {
 		HandshakeInfo handshakeInfo = session.getHandshakeInfo();
@@ -145,62 +112,49 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 								Mono.empty()))
 				.subscribe();
 
-		return session.send(session.receive().flatMap((message) -> {
-			Map<String, Object> map = decode(message);
-			String id = (String) map.get("id");
-			MessageType messageType = MessageType.resolve((String) map.get("type"));
-			if (messageType == null) {
-				return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
-			}
-			switch (messageType) {
-			case SUBSCRIBE:
-				if (!connectionInitProcessed.get()) {
-					return GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS);
-				}
-				if (id == null) {
-					return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
-				}
-				WebInput input = new WebInput(
-						handshakeInfo.getUri(), handshakeInfo.getHeaders(), getPayload(map), null, id);
-				if (logger.isDebugEnabled()) {
-					logger.debug("Executing: " + input);
-				}
-				return this.graphQlHandler.handleRequest(input)
-						.flatMapMany((output) -> handleWebOutput(session, id, subscriptions, output))
-						.doOnTerminate(() -> subscriptions.remove(id));
-			case COMPLETE:
-				if (id != null) {
-					Subscription subscription = subscriptions.remove(id);
-					if (subscription != null) {
-						subscription.cancel();
+		return session.send(session.receive().flatMap(webSocketMessage -> {
+			GraphQlWebSocketMessage message = this.codecDelegate.decode(webSocketMessage);
+			String id = message.getId();
+			Map<String, Object> payload = message.getPayloadOrDefault(Collections.emptyMap());
+			switch (message.getType()) {
+				case "subscribe":
+					if (!connectionInitProcessed.get()) {
+						return GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS);
 					}
-				}
-				return this.graphQlHandler.handleWebSocketCompletion().thenMany(Flux.empty());
-			case CONNECTION_INIT:
-				if (!connectionInitProcessed.compareAndSet(false, true)) {
-					return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
-				}
-				return this.graphQlHandler.handleWebSocketInitialization(getPayload(map))
-						.defaultIfEmpty(Collections.emptyMap())
-						.flatMapMany(ackPayload -> Flux.just(encode(session, null, MessageType.CONNECTION_ACK, ackPayload)))
-						.onErrorResume(ex -> GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS));
-			default:
-				return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
+					if (id == null) {
+						return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
+					}
+					WebInput input = new WebInput(
+							handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload, null, id);
+					if (logger.isDebugEnabled()) {
+						logger.debug("Executing: " + input);
+					}
+					return this.graphQlHandler.handleRequest(input)
+							.flatMapMany((output) -> handleWebOutput(session, id, subscriptions, output))
+							.doOnTerminate(() -> subscriptions.remove(id));
+				case "complete":
+					if (id != null) {
+						Subscription subscription = subscriptions.remove(id);
+						if (subscription != null) {
+							subscription.cancel();
+						}
+					}
+					return this.graphQlHandler.handleWebSocketCompletion().thenMany(Flux.empty());
+				case "connection_init":
+					if (!connectionInitProcessed.compareAndSet(false, true)) {
+						return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
+					}
+					return this.graphQlHandler.handleWebSocketInitialization(payload)
+							.defaultIfEmpty(Collections.emptyMap())
+							.map(ackPayload -> this.codecDelegate.encodeConnectionAckMessage(session, ackPayload))
+							.flux()
+							.onErrorResume(ex -> GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS));
+				default:
+					return GraphQlStatus.close(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
 			}
 		}));
 	}
 
-	@SuppressWarnings({ "unchecked", "ConstantConditions" })
-	private Map<String, Object> decode(WebSocketMessage message) {
-		DataBuffer buffer = DataBufferUtils.retain(message.getPayload());
-		return (Map<String, Object>) this.decoder.decode(buffer, MAP_RESOLVABLE_TYPE, null, null);
-	}
-
-	@SuppressWarnings("unchecked")
-	private static Map<String, Object> getPayload(Map<String, Object> message) {
-		Map<String, Object> payload = (Map<String, Object>) message.get("payload");
-		return (payload != null ? payload : Collections.emptyMap());
-	}
 
 	@SuppressWarnings("unchecked")
 	private Flux<WebSocketMessage> handleWebOutput(WebSocketSession session, String id,
@@ -230,78 +184,17 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 		}
 
 		return outputFlux
-				.map((result) -> {
-						Map<String, Object> dataMap = result.toSpecification();
-						return encode(session, id, MessageType.NEXT, dataMap);
-				})
-				.concatWith(Mono.fromCallable(() -> encode(session, id, MessageType.COMPLETE, null)))
-				.onErrorResume((ex) -> {
+				.map(result -> this.codecDelegate.encodeNextMessage(session, id, result))
+				.concatWith(Mono.fromCallable(() -> this.codecDelegate.encodeCompleteMessage(session, id)))
+				.onErrorResume(ex -> {
 						if (ex instanceof SubscriptionExistsException) {
 							CloseStatus status = new CloseStatus(4409, "Subscriber for " + id + " already exists");
 							return GraphQlStatus.close(session, status);
 						}
-						Map<String, Object> errorMap = GraphqlErrorBuilder.newError()
-								.errorType(ErrorType.DataFetchingException)
-								.message(ex.getMessage())
-								.build()
-								.toSpecification();
-						return Mono.just(encode(
-								session, id, MessageType.ERROR, Collections.singletonList(errorMap)));
+						return Mono.fromCallable(() -> this.codecDelegate.encodeErrorMessage(session, id, ex));
 				});
 	}
 
-	@SuppressWarnings("unchecked")
-	private <T> WebSocketMessage encode(WebSocketSession session, @Nullable String id, MessageType messageType,
-			@Nullable Object payload) {
-
-		Map<String, Object> payloadMap = new HashMap<>(3);
-		if (id != null) {
-			payloadMap.put("id", id);
-		}
-		payloadMap.put("type", messageType.getType());
-		if (payload != null) {
-			payloadMap.put("payload", payload);
-		}
-
-		DataBuffer buffer = ((Encoder<T>) this.encoder).encodeValue((T) payloadMap, session.bufferFactory(),
-				MAP_RESOLVABLE_TYPE, MimeTypeUtils.APPLICATION_JSON, null);
-
-		return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer);
-	}
-
-	private enum MessageType {
-
-		CONNECTION_INIT("connection_init"),
-		CONNECTION_ACK("connection_ack"),
-		SUBSCRIBE("subscribe"),
-		NEXT("next"),
-		ERROR("error"),
-		COMPLETE("complete");
-
-		private static final Map<String, MessageType> messageTypes = new HashMap<>(6);
-
-		static {
-			for (MessageType messageType : MessageType.values()) {
-				messageTypes.put(messageType.getType(), messageType);
-			}
-		}
-
-		private final String type;
-
-		MessageType(String type) {
-			this.type = type;
-		}
-
-		public String getType() {
-			return this.type;
-		}
-
-		@Nullable
-		public static MessageType resolve(@Nullable String type) {
-			return (type != null) ? messageTypes.get(type) : null;
-		}
-
-	}
 
 	private static class GraphQlStatus {
 
@@ -319,9 +212,9 @@ public class GraphQlWebSocketHandler implements WebSocketHandler {
 
 	}
 
+
 	@SuppressWarnings("serial")
 	private static class SubscriptionExistsException extends RuntimeException {
-
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketMessage.java

@@ -0,0 +1,165 @@
+/*
+ * 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.web.webflux;
+
+import java.util.Collections;
+
+import graphql.ExecutionResult;
+import graphql.GraphQLError;
+
+import org.springframework.graphql.RequestInput;
+import org.springframework.lang.Nullable;
+import org.springframework.util.ObjectUtils;
+
+/**
+ * Representation of a GraphQL over WebSocket protocol message.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public class GraphQlWebSocketMessage {
+
+	@Nullable
+	private String id;
+
+	private String type;
+
+	@Nullable
+	private Object payload;
+
+
+	/**
+	 * Private constructor for static factory methods.
+	 */
+	private GraphQlWebSocketMessage(@Nullable String id, String type, @Nullable Object payload) {
+		this.id = id;
+		this.type = type;
+		this.payload = payload;
+	}
+
+	/**
+	 * Constructor for deserialization.
+	 */
+	GraphQlWebSocketMessage() {
+		this.type = "";
+	}
+
+
+	@Nullable
+	public String getId() {
+		return this.id;
+	}
+
+	public String getType() {
+		return this.type;
+	}
+
+	@SuppressWarnings("unchecked")
+	@Nullable
+	public <P> P getPayload() {
+		return (P) this.payload;
+	}
+
+	@SuppressWarnings("unchecked")
+	public <P> P getPayloadOrDefault(P defaultPayload) {
+		return (this.payload != null ? (P) this.payload : defaultPayload);
+	}
+
+	public void setId(@Nullable String id) {
+		this.id = id;
+	}
+
+	public void setType(String type) {
+		this.type = type;
+	}
+
+	public void setPayload(@Nullable Object payload) {
+		this.payload = payload;
+	}
+
+
+	@Override
+	public int hashCode() {
+		int hashCode = this.type.hashCode();
+		hashCode = 31 * hashCode + ObjectUtils.nullSafeHashCode(this.id);
+		hashCode = 31 * hashCode + ObjectUtils.nullSafeHashCode(this.payload);
+		return hashCode;
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		if (!(o instanceof GraphQlWebSocketMessage)) {
+			return false;
+		}
+		GraphQlWebSocketMessage other = (GraphQlWebSocketMessage) o;
+		return (this.type.equals(other.type) &&
+				(ObjectUtils.nullSafeEquals(this.id, other.id) || (this.id == null && other.id == null)) &&
+				(ObjectUtils.nullSafeEquals(this.payload, other.payload) || (this.payload == null && other.payload == null)));
+	}
+
+	@Override
+	public String toString() {
+		return "GraphQlWebSocketMessage[" +
+				(this.id != null ? "id=\"" + this.id + "\"" + ", " : "") +
+				"type=\"" + this.type + "\"" +
+				(this.payload != null ? ", payload=" + this.payload : "") + "]";
+	}
+
+
+	/**
+	 * Create a "connection_init" message.
+	 */
+	public static GraphQlWebSocketMessage connectionInit(@Nullable Object payload) {
+		return new GraphQlWebSocketMessage(null, "connection_init", payload);
+	}
+
+	/**
+	 * Create a "connection_ack" message.
+	 */
+	public static GraphQlWebSocketMessage connectionAck(@Nullable Object payload) {
+		return new GraphQlWebSocketMessage(null, "connection_ack", payload);
+	}
+
+	/**
+	 * Create a "subscribe" message.
+	 */
+	public static GraphQlWebSocketMessage subscribe(String id, RequestInput input) {
+		return new GraphQlWebSocketMessage(id, "subscribe", input.toMap());
+	}
+
+	/**
+	 * Create a "next" message.
+	 */
+	public static GraphQlWebSocketMessage next(String id, ExecutionResult result) {
+		return new GraphQlWebSocketMessage(id, "next", result.toSpecification());
+	}
+
+	/**
+	 * Create an "error" message.
+	 */
+	public static GraphQlWebSocketMessage error(String id, GraphQLError error) {
+		return new GraphQlWebSocketMessage(id, "error", Collections.singletonList(error.toSpecification()));
+	}
+
+	/**
+	 * Create a "complete" message.
+	 */
+	public static GraphQlWebSocketMessage complete(String id) {
+		return new GraphQlWebSocketMessage(id, "complete", null);
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/web/webflux/WebSocketCodecDelegate.java

@@ -0,0 +1,107 @@
+/*
+ * 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.web.webflux;
+
+import graphql.ExecutionResult;
+import graphql.GraphQLError;
+import graphql.GraphqlErrorBuilder;
+
+import org.springframework.core.ResolvableType;
+import org.springframework.core.codec.Decoder;
+import org.springframework.core.codec.Encoder;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
+import org.springframework.http.MediaType;
+import org.springframework.http.codec.CodecConfigurer;
+import org.springframework.http.codec.DecoderHttpMessageReader;
+import org.springframework.http.codec.EncoderHttpMessageWriter;
+import org.springframework.util.Assert;
+import org.springframework.util.MimeTypeUtils;
+import org.springframework.web.reactive.socket.WebSocketMessage;
+import org.springframework.web.reactive.socket.WebSocketSession;
+
+/**
+ * WebFlux Support class for GraphQL over WebSocket handling.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+final class WebSocketCodecDelegate {
+
+	private static final ResolvableType MESSAGE_TYPE = ResolvableType.forClass(GraphQlWebSocketMessage.class);
+
+
+	private final Decoder<?> decoder;
+
+	private final Encoder<?> encoder;
+
+
+	WebSocketCodecDelegate(CodecConfigurer codecConfigurer) {
+		Assert.notNull(codecConfigurer, "CodecConfigurer is required");
+		this.decoder = initDecoder(codecConfigurer);
+		this.encoder = initEncoder(codecConfigurer);
+	}
+
+	private static Decoder<?> initDecoder(CodecConfigurer configurer) {
+		return configurer.getReaders().stream()
+				.filter((reader) -> reader.canRead(MESSAGE_TYPE, MediaType.APPLICATION_JSON))
+				.map((reader) -> ((DecoderHttpMessageReader<?>) reader).getDecoder())
+				.findFirst()
+				.orElseThrow(() -> new IllegalArgumentException("No JSON Decoder"));
+	}
+
+	private static Encoder<?> initEncoder(CodecConfigurer configurer) {
+		return configurer.getWriters().stream()
+				.filter((writer) -> writer.canWrite(MESSAGE_TYPE, MediaType.APPLICATION_JSON))
+				.map((writer) -> ((EncoderHttpMessageWriter<?>) writer).getEncoder())
+				.findFirst()
+				.orElseThrow(() -> new IllegalArgumentException("No JSON Encoder"));
+	}
+
+
+	@SuppressWarnings("ConstantConditions")
+	public GraphQlWebSocketMessage decode(WebSocketMessage webSocketMessage) {
+		DataBuffer buffer = DataBufferUtils.retain(webSocketMessage.getPayload());
+		return (GraphQlWebSocketMessage) this.decoder.decode(buffer, MESSAGE_TYPE, null, null);
+	}
+
+	public WebSocketMessage encodeConnectionAckMessage(WebSocketSession session, Object ackPayload) {
+		return encode(session, GraphQlWebSocketMessage.connectionAck(ackPayload));
+	}
+
+	public WebSocketMessage encodeNextMessage(WebSocketSession session, String id, ExecutionResult result) {
+		return encode(session, GraphQlWebSocketMessage.next(id, result));
+	}
+
+	public WebSocketMessage encodeErrorMessage(WebSocketSession session, String id, Throwable ex) {
+		GraphQLError error = GraphqlErrorBuilder.newError().message(ex.getMessage()).build();
+		return encode(session, GraphQlWebSocketMessage.error(id, error));
+	}
+
+	public WebSocketMessage encodeCompleteMessage(WebSocketSession session, String id) {
+		return encode(session, GraphQlWebSocketMessage.complete(id));
+	}
+
+	@SuppressWarnings("unchecked")
+	private  <T> WebSocketMessage encode(WebSocketSession session, GraphQlWebSocketMessage message) {
+
+		DataBuffer buffer = ((Encoder<T>) this.encoder).encodeValue(
+				(T) message, session.bufferFactory(), MESSAGE_TYPE, MimeTypeUtils.APPLICATION_JSON, null);
+
+		return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer);
+	}
+
+}

spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java

@@ -25,13 +25,12 @@ import java.net.URI;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
-import graphql.ErrorType;
 import graphql.ExecutionResult;
+import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -46,11 +45,12 @@ import reactor.core.scheduler.Schedulers;
 import org.springframework.graphql.web.WebGraphQlHandler;
 import org.springframework.graphql.web.WebInput;
 import org.springframework.graphql.web.WebOutput;
+import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpInputMessage;
 import org.springframework.http.HttpOutputMessage;
+import org.springframework.http.converter.GenericHttpMessageConverter;
 import org.springframework.http.converter.HttpMessageConverter;
-import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.web.socket.CloseStatus;
@@ -132,18 +132,13 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 	}
 
 	@Override
-	@SuppressWarnings("unchecked")
-	protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
-		Map<String, Object> map = decode(message, Map.class);
-		String id = (String) map.get("id");
-		MessageType messageType = MessageType.resolve((String) map.get("type"));
-		if (messageType == null) {
-			GraphQlStatus.closeSession(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
-			return;
-		}
+	protected void handleTextMessage(WebSocketSession session, TextMessage webSocketMessage) throws Exception {
+		GraphQlWebSocketMessage message = decode(webSocketMessage);
+		String id = message.getId();
+		Map<String, Object> payload = message.getPayloadOrDefault(Collections.emptyMap());
 		SessionState sessionState = getSessionInfo(session);
-		switch (messageType) {
-		case SUBSCRIBE:
+		switch (message.getType()) {
+		case "subscribe":
 			if (sessionState.isConnectionInitNotProcessed()) {
 				GraphQlStatus.closeSession(session, GraphQlStatus.UNAUTHORIZED_STATUS);
 				return;
@@ -155,7 +150,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 			URI uri = session.getUri();
 			Assert.notNull(uri, "Expected handshake url");
 			HttpHeaders headers = session.getHandshakeHeaders();
-			WebInput input = new WebInput(uri, headers, getPayload(map), null, id);
+			WebInput input = new WebInput(uri, headers, payload, null, id);
 			if (logger.isDebugEnabled()) {
 				logger.debug("Executing: " + input);
 			}
@@ -164,7 +159,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 					.publishOn(sessionState.getScheduler()) // Serial blocking send via single thread
 					.subscribe(new SendMessageSubscriber(id, session, sessionState));
 			return;
-		case COMPLETE:
+		case "complete":
 			if (id != null) {
 				Subscription subscription = sessionState.getSubscriptions().remove(id);
 				if (subscription != null) {
@@ -173,16 +168,16 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 			}
 			this.graphQlHandler.handleWebSocketCompletion().block(Duration.ofSeconds(10));
 			return;
-		case CONNECTION_INIT:
+		case "connection_init":
 			if (sessionState.setConnectionInitProcessed()) {
 				GraphQlStatus.closeSession(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS);
 				return;
 			}
-			this.graphQlHandler.handleWebSocketInitialization(getPayload(map))
+			this.graphQlHandler.handleWebSocketInitialization(payload)
 					.defaultIfEmpty(Collections.emptyMap())
 					.publishOn(sessionState.getScheduler()) // Serial blocking send via single thread
 					.doOnNext(ackPayload -> {
-						TextMessage outputMessage = encode(null, MessageType.CONNECTION_ACK, ackPayload);
+						TextMessage outputMessage = encode(GraphQlWebSocketMessage.connectionAck(ackPayload));
 						try {
 							session.sendMessage(outputMessage);
 						}
@@ -199,18 +194,12 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 		default:
 			GraphQlStatus.closeSession(session, GraphQlStatus.INVALID_MESSAGE_STATUS);
 		}
-
 	}
 
 	@SuppressWarnings("unchecked")
-	private <T> T decode(TextMessage message, Class<T> targetClass) throws IOException {
-		return ((HttpMessageConverter<T>) this.converter).read(targetClass, new HttpInputMessageAdapter(message));
-	}
-
-	@SuppressWarnings("unchecked")
-	private static Map<String, Object> getPayload(Map<String, Object> message) {
-		Map<String, Object> payload = (Map<String, Object>) message.get("payload");
-		return (payload != null ? payload : Collections.emptyMap());
+	private GraphQlWebSocketMessage decode(TextMessage message) throws IOException {
+		return ((GenericHttpMessageConverter<GraphQlWebSocketMessage>) this.converter)
+				.read(GraphQlWebSocketMessage.class, null, new HttpInputMessageAdapter(message));
 	}
 
 	private SessionState getSessionInfo(WebSocketSession session) {
@@ -244,46 +233,29 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 		}
 
 		return outputFlux
-				.map((result) -> {
-						Map<String, Object> dataMap = result.toSpecification();
-						return encode(id, MessageType.NEXT, dataMap);
-				})
-				.concatWith(Mono.fromCallable(() -> encode(id, MessageType.COMPLETE, null)))
+				.map(result -> encode(GraphQlWebSocketMessage.next(id, result)))
+				.concatWith(Mono.fromCallable(() -> encode(GraphQlWebSocketMessage.complete(id))))
 				.onErrorResume((ex) -> {
 						if (ex instanceof SubscriptionExistsException) {
 							CloseStatus status = new CloseStatus(4409, "Subscriber for " + id + " already exists");
 							GraphQlStatus.closeSession(session, status);
 							return Flux.empty();
 						}
-						ErrorType errorType = ErrorType.DataFetchingException;
 						String message = ex.getMessage();
-						Map<String, Object> errorMap = GraphqlErrorBuilder.newError()
-								.errorType(errorType)
-								.message(message)
-								.build()
-								.toSpecification();
-						return Mono.just(encode(
-								id, MessageType.ERROR, Collections.singletonList(errorMap)));
+						GraphQLError error = GraphqlErrorBuilder.newError().message(message).build();
+						return Mono.just(encode(GraphQlWebSocketMessage.error(id, error)));
 				});
 	}
 
 	@SuppressWarnings("unchecked")
-	private <T> TextMessage encode(@Nullable String id, MessageType messageType, @Nullable Object payload) {
-		Map<String, Object> payloadMap = new HashMap<>(3);
-		payloadMap.put("type", messageType.getType());
-		if (id != null) {
-			payloadMap.put("id", id);
-		}
-		if (payload != null) {
-			payloadMap.put("payload", payload);
-		}
+	private <T> TextMessage encode(GraphQlWebSocketMessage message) {
 		try {
 			HttpOutputMessageAdapter outputMessage = new HttpOutputMessageAdapter();
-			((HttpMessageConverter<T>) this.converter).write((T) payloadMap, null, outputMessage);
+			((HttpMessageConverter<T>) this.converter).write((T) message, null, outputMessage);
 			return new TextMessage(outputMessage.toByteArray());
 		}
 		catch (IOException ex) {
-			throw new IllegalStateException("Failed to write " + payloadMap + " as JSON", ex);
+			throw new IllegalStateException("Failed to write " + message + " as JSON", ex);
 		}
 	}
 
@@ -308,39 +280,6 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 		return false;
 	}
 
-	private enum MessageType {
-
-		CONNECTION_INIT("connection_init"),
-		CONNECTION_ACK("connection_ack"),
-		SUBSCRIBE("subscribe"),
-		NEXT("next"),
-		ERROR("error"),
-		COMPLETE("complete");
-
-		private static final Map<String, MessageType> messageTypes = new HashMap<>(6);
-
-		static {
-			for (MessageType messageType : MessageType.values()) {
-				messageTypes.put(messageType.getType(), messageType);
-			}
-		}
-
-		private final String type;
-
-		MessageType(String type) {
-			this.type = type;
-		}
-
-		public String getType() {
-			return this.type;
-		}
-
-		@Nullable
-		public static MessageType resolve(@Nullable String type) {
-			return (type != null) ? messageTypes.get(type) : null;
-		}
-
-	}
 
 	private static class GraphQlStatus {
 
@@ -488,7 +427,6 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 
 	@SuppressWarnings("serial")
 	private static class SubscriptionExistsException extends RuntimeException {
-
 	}
 
 }

spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java

@@ -26,19 +26,20 @@ import java.util.function.BiConsumer;
 
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
-import org.springframework.graphql.GraphQlSetup;
-import org.springframework.graphql.web.WebGraphQlHandler;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
 
+import org.springframework.core.ResolvableType;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
-import org.springframework.graphql.web.WebSocketHandlerTestSupport;
+import org.springframework.graphql.GraphQlSetup;
 import org.springframework.graphql.web.ConsumeOneAndNeverCompleteInterceptor;
+import org.springframework.graphql.web.WebGraphQlHandler;
 import org.springframework.graphql.web.WebInterceptor;
+import org.springframework.graphql.web.WebSocketHandlerTestSupport;
 import org.springframework.graphql.web.WebSocketInterceptor;
 import org.springframework.http.codec.ServerCodecConfigurer;
 import org.springframework.http.codec.json.Jackson2JsonDecoder;
@@ -63,12 +64,15 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 		StepVerifier.create(session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
-				.consumeNextWith((message) -> assertThat(decode(message)).hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID).containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("bookById", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("name", "Nineteen Eighty-Four"))
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("next");
+					assertThat(actual.<Map<String, Object>>getPayload())
+							.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.extractingByKey("bookById", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.containsEntry("name", "Nineteen Eighty-Four");
+				})
 				.consumeNextWith((message) -> assertMessageType(message, "complete"))
 				.verifyComplete();
 	}
@@ -79,13 +83,15 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 				toWebSocketMessage("{\"type\":\"connection_init\"}"),
 				toWebSocketMessage(BOOK_SUBSCRIPTION)));
 
-		BiConsumer<WebSocketMessage, String> bookPayloadAssertion = (message, bookId) ->
-				assertThat(decode(message))
-						.hasSize(3).containsEntry("id", SUBSCRIPTION_ID).containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("bookSearch", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("id", bookId);
+		BiConsumer<WebSocketMessage, String> bookPayloadAssertion = (message, bookId) -> {
+			GraphQlWebSocketMessage actual = decode(message);
+			assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+			assertThat(actual.getType()).isEqualTo("next");
+			assertThat(actual.<Map<String, Object>>getPayload())
+					.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.extractingByKey("bookSearch", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.containsEntry("id", bookId);
+		};
 
 		StepVerifier.create(session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
@@ -128,7 +134,6 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 	}
 
 	@Test
-	@SuppressWarnings("unchecked")
 	void connectionInitHandling() {
 		TestWebSocketSession session = handle(
 				Flux.just(toWebSocketMessage("{\"type\":\"connection_init\",\"payload\":{\"key\":\"A\"}}")),
@@ -143,9 +148,9 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 		StepVerifier.create(session.getOutput())
 				.consumeNextWith((message) -> {
-					Map<String, Object> content = decode(message);
-					assertThat(content).containsEntry("type", "connection_ack");
-					assertThat((Map<String, Object>) content.get("payload")).containsEntry("key", "A acknowledged");
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getType()).isEqualTo("connection_ack");
+					assertThat(actual.<Map<String, Object>>getPayload()).containsEntry("key", "A acknowledged");
 				})
 				.verifyComplete();
 	}
@@ -214,7 +219,7 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 		TestWebSocketSession session = handle(messageFlux, new ConsumeOneAndNeverCompleteInterceptor());
 
 		// Collect messages until session closed
-		List<Map<String, Object>> messages = new ArrayList<>();
+		List<GraphQlWebSocketMessage> messages = new ArrayList<>();
 		session.getOutput().subscribe((message) -> messages.add(decode(message)));
 
 		StepVerifier.create(session.closeStatus())
@@ -222,8 +227,8 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 				.verifyComplete();
 
 		assertThat(messages.size()).isEqualTo(2);
-		assertThat(messages.get(0).get("type")).isEqualTo("connection_ack");
-		assertThat(messages.get(1).get("type")).isEqualTo("next");
+		assertThat(messages.get(0).getType()).isEqualTo("connection_ack");
+		assertThat(messages.get(1).getType()).isEqualTo("next");
 	}
 
 	@Test
@@ -277,27 +282,28 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 		StepVerifier.create(session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
-				.consumeNextWith((message) -> assertThat(decode(message))
-						.hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID)
-						.containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("greeting", "a"))
-				.consumeNextWith((message) -> assertThat(decode(message))
-						.hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID)
-						.containsEntry("type", "error")
-						.hasEntrySatisfying("payload", payload -> assertThat(payload)
-								.asList()
-								.hasSize(1)
-								.allSatisfy(theError -> assertThat(theError)
-										.asInstanceOf(InstanceOfAssertFactories.map(String.class, Object.class))
-										.hasSize(3)
-										.hasEntrySatisfying("locations", loc -> assertThat(loc).asList().isEmpty())
-										.hasEntrySatisfying("message", msg -> assertThat(msg).asString().contains("null"))
-										.extractingByKey("extensions", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-										.containsEntry("classification", "DataFetchingException"))))
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("next");
+					assertThat(actual.<Map<String, Object>>getPayload())
+							.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.containsEntry("greeting", "a");
+				})
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("error");
+					assertThat(actual.<List<Map<String, Object>>>getPayload())
+							.asList().hasSize(1)
+							.allSatisfy(theError -> assertThat(theError)
+									.asInstanceOf(InstanceOfAssertFactories.map(String.class, Object.class))
+									.hasSize(3)
+									.hasEntrySatisfying("locations", loc -> assertThat(loc).asList().isEmpty())
+									.hasEntrySatisfying("message", msg -> assertThat(msg).asString().contains("null"))
+									.extractingByKey("extensions", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+									.containsEntry("classification", "DataFetchingException"));
+				})
 				.verifyComplete();
 	}
 
@@ -317,17 +323,17 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 		return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer);
 	}
 
-	@SuppressWarnings({ "unchecked", "ConstantConditions" })
-	private Map<String, Object> decode(WebSocketMessage message) {
-		return (Map<String, Object>) decoder.decode(DataBufferUtils.retain(message.getPayload()),
-				GraphQlWebSocketHandler.MAP_RESOLVABLE_TYPE, null, Collections.emptyMap());
+	@SuppressWarnings("ConstantConditions")
+	private GraphQlWebSocketMessage decode(WebSocketMessage message) {
+		return (GraphQlWebSocketMessage) decoder.decode(DataBufferUtils.retain(message.getPayload()),
+				ResolvableType.forClass(GraphQlWebSocketMessage.class), null, Collections.emptyMap());
 	}
 
-	private void assertMessageType(WebSocketMessage message, String messageType) {
-		Map<String, Object> map = decode(message);
-		assertThat(map).containsEntry("type", messageType);
+	private void assertMessageType(WebSocketMessage webSocketMessage, String messageType) {
+		GraphQlWebSocketMessage message = decode(webSocketMessage);
+		assertThat(message.getType()).isEqualTo(messageType);
 		if (!messageType.equals("connection_ack")) {
-			assertThat(map).containsEntry("id", SUBSCRIPTION_ID);
+			assertThat(message.getId()).isEqualTo(SUBSCRIPTION_ID);
 		}
 	}
 

spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java

@@ -29,18 +29,20 @@ import java.util.function.Consumer;
 
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
-import org.springframework.graphql.GraphQlSetup;
-import org.springframework.graphql.web.WebGraphQlHandler;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
-import org.springframework.graphql.web.WebSocketHandlerTestSupport;
+import org.springframework.graphql.GraphQlSetup;
 import org.springframework.graphql.web.ConsumeOneAndNeverCompleteInterceptor;
+import org.springframework.graphql.web.WebGraphQlHandler;
 import org.springframework.graphql.web.WebInterceptor;
+import org.springframework.graphql.web.WebSocketHandlerTestSupport;
 import org.springframework.graphql.web.WebSocketInterceptor;
+import org.springframework.graphql.web.webflux.GraphQlWebSocketMessage;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpInputMessage;
+import org.springframework.http.converter.GenericHttpMessageConverter;
 import org.springframework.http.converter.HttpMessageConverter;
 import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter;
 import org.springframework.web.socket.CloseStatus;
@@ -69,12 +71,15 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 		StepVerifier.create(this.session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
-				.consumeNextWith((message) -> assertThat(decode(message)).hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID).containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("bookById", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("name", "Nineteen Eighty-Four"))
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("next");
+					assertThat(actual.<Map<String, Object>>getPayload())
+							.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.extractingByKey("bookById", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.containsEntry("name", "Nineteen Eighty-Four");
+				})
 				.consumeNextWith((message) -> assertMessageType(message, "complete"))
 				.then(this.session::close) // Complete output Flux
 				.verifyComplete();
@@ -84,13 +89,15 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 	void subscription() throws Exception {
 		handle(this.handler, new TextMessage("{\"type\":\"connection_init\"}"), new TextMessage(BOOK_SUBSCRIPTION));
 
-		BiConsumer<WebSocketMessage<?>, String> bookPayloadAssertion = (message, bookId) ->
-				assertThat(decode(message))
-						.hasSize(3).containsEntry("id", SUBSCRIPTION_ID).containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("bookSearch", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("id", bookId);
+		BiConsumer<WebSocketMessage<?>, String> bookPayloadAssertion = (message, bookId) -> {
+			GraphQlWebSocketMessage actual = decode(message);
+			assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+			assertThat(actual.getType()).isEqualTo("next");
+			assertThat(actual.<Map<String, Object>>getPayload())
+					.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.extractingByKey("bookSearch", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+					.containsEntry("id", bookId);
+		};
 
 		StepVerifier.create(this.session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
@@ -129,7 +136,6 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 	}
 
 	@Test
-	@SuppressWarnings("unchecked")
 	void connectionInitHandling() throws Exception {
 
 		WebSocketInterceptor interceptor = new WebSocketInterceptor() {
@@ -145,10 +151,10 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 				new TextMessage("{\"type\":\"connection_init\",\"payload\":{\"key\":\"A\"}}"));
 
 		StepVerifier.create(session.getOutput())
-				.consumeNextWith((message) -> {
-					Map<String, Object> content = decode(message);
-					assertThat(content).containsEntry("type", "connection_ack");
-					assertThat((Map<String, Object>) content.get("payload")).containsEntry("key", "A acknowledged");
+				.consumeNextWith((webSocketMessage) -> {
+					GraphQlWebSocketMessage message = decode(webSocketMessage);
+					assertThat(message.getType()).isEqualTo("connection_ack");
+					assertThat(message.<Map<String, Object>>getPayload()).containsEntry("key", "A acknowledged");
 				})
 				.then(this.session::close) // Complete output Flux
 				.verifyComplete();
@@ -211,7 +217,7 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 				new TextMessage(BOOK_SUBSCRIPTION));
 
 		// Collect messages until session closed
-		List<Map<String, Object>> messages = new ArrayList<>();
+		List<GraphQlWebSocketMessage> messages = new ArrayList<>();
 		this.session.getOutput().subscribe((message) -> messages.add(decode(message)));
 
 		StepVerifier.create(this.session.closeStatus())
@@ -219,8 +225,8 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 				.verifyComplete();
 
 		assertThat(messages.size()).isEqualTo(2);
-		assertThat(messages.get(0).get("type")).isEqualTo("connection_ack");
-		assertThat(messages.get(1).get("type")).isEqualTo("next");
+		assertThat(messages.get(0).getType()).isEqualTo("connection_ack");
+		assertThat(messages.get(1).getType()).isEqualTo("next");
 	}
 
 	@Test
@@ -278,27 +284,28 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 
 		StepVerifier.create(this.session.getOutput())
 				.consumeNextWith((message) -> assertMessageType(message, "connection_ack"))
-				.consumeNextWith((message) -> assertThat(decode(message))
-						.hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID)
-						.containsEntry("type", "next")
-						.extractingByKey("payload", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-						.containsEntry("greeting", "a"))
-				.consumeNextWith((message) -> assertThat(decode(message))
-						.hasSize(3)
-						.containsEntry("id", SUBSCRIPTION_ID)
-						.containsEntry("type", "error")
-						.hasEntrySatisfying("payload", payload -> assertThat(payload)
-								.asList()
-								.hasSize(1)
-								.allSatisfy(theError -> assertThat(theError)
-										.asInstanceOf(InstanceOfAssertFactories.map(String.class, Object.class))
-										.hasSize(3)
-										.hasEntrySatisfying("locations", loc -> assertThat(loc).asList().isEmpty())
-										.hasEntrySatisfying("message", msg -> assertThat(msg).asString().contains("null"))
-										.extractingByKey("extensions", as(InstanceOfAssertFactories.map(String.class, Object.class)))
-										.containsEntry("classification", "DataFetchingException"))))
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("next");
+					assertThat(actual.<Map<String, Object>>getPayload())
+							.extractingByKey("data", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+							.containsEntry("greeting", "a");
+				})
+				.consumeNextWith((message) -> {
+					GraphQlWebSocketMessage actual = decode(message);
+					assertThat(actual.getId()).isEqualTo(SUBSCRIPTION_ID);
+					assertThat(actual.getType()).isEqualTo("error");
+					assertThat(actual.<List<Map<String, Object>>>getPayload())
+							.asList().hasSize(1)
+							.allSatisfy(theError -> assertThat(theError)
+									.asInstanceOf(InstanceOfAssertFactories.map(String.class, Object.class))
+									.hasSize(3)
+									.hasEntrySatisfying("locations", loc -> assertThat(loc).asList().isEmpty())
+									.hasEntrySatisfying("message", msg -> assertThat(msg).asString().contains("null"))
+									.extractingByKey("extensions", as(InstanceOfAssertFactories.map(String.class, Object.class)))
+									.containsEntry("classification", "DataFetchingException"));
+				})
 				.then(this.session::close)
 				.verifyComplete();
 	}
@@ -319,30 +326,27 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport {
 		}
 	}
 
-	private void assertMessageType(WebSocketMessage<?> message, String messageType) {
-		Map<String, Object> map = decode(message, Map.class);
-		assertThat(map).containsEntry("type", messageType);
-		if (!messageType.equals("connection_ack")) {
-			assertThat(map).containsEntry("id", SUBSCRIPTION_ID);
-		}
-	}
-
 	@SuppressWarnings("unchecked")
-	private Map<String, Object> decode(WebSocketMessage<?> message) {
-		return decode(message, Map.class);
-	}
-
-	@SuppressWarnings("unchecked")
-	private <T> T decode(WebSocketMessage<?> message, Class<T> targetClass) {
+	private GraphQlWebSocketMessage decode(WebSocketMessage<?> message) {
 		try {
 			HttpInputMessageAdapter inputMessage = new HttpInputMessageAdapter((TextMessage) message);
-			return ((HttpMessageConverter<T>) converter).read(targetClass, inputMessage);
+			return ((GenericHttpMessageConverter<GraphQlWebSocketMessage>) converter)
+					.read(GraphQlWebSocketMessage.class, null, inputMessage);
 		}
 		catch (IOException ex) {
 			throw new IllegalStateException(ex);
 		}
 	}
 
+	private void assertMessageType(WebSocketMessage<?> webSocketMessage, String messageType) {
+		GraphQlWebSocketMessage message = decode(webSocketMessage);
+		assertThat(message.getType()).isEqualTo(messageType);
+		if (!messageType.equals("connection_ack")) {
+			assertThat(message.getId()).isEqualTo(SUBSCRIPTION_ID);
+		}
+	}
+
+
 	private static class HttpInputMessageAdapter extends ByteArrayInputStream implements HttpInputMessage {
 
 		HttpInputMessageAdapter(TextMessage message) {