From dbba8b17f8a52623ee7e7f41306c3d4ff81a0a49 Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Tue, 12 Apr 2022 21:23:02 +0100
Subject: [PATCH] Uncomment assertion for context propagation

Closes gh-342
---
 .../graphql/server/webmvc/GraphQlWebSocketHandler.java   | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
index f4f397f2..588be8e5 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java
@@ -365,12 +365,9 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub
 		public Closeable restoreThreadLocalValue(WebSocketSession session) {
 			if (this.accessor != null) {
 				Map<String, Object> valuesMap = (Map<String, Object>) session.getAttributes().get(SAVED_CONTEXT_KEY);
-				// Uncomment when Boot is updated to use HandshakeInterceptor
-				// Assert.state(valuesMap != null, "No context");
-				if (valuesMap != null) {
-					this.accessor.restoreValues(valuesMap);
-					return () -> this.accessor.resetValues(valuesMap);
-				}
+				Assert.state(valuesMap != null, "No ThreadLocal context in WebSocketSession attributes");
+				this.accessor.restoreValues(valuesMap);
+				return () -> this.accessor.resetValues(valuesMap);
 			}
 			return () -> {};
 		}