From 1289dbae84d7dc0e07bf6701e6ee0f4441783f67 Mon Sep 17 00:00:00 2001
From: Athou <jeremiepanzer@gmail.com>
Date: Thu, 28 Dec 2023 10:25:44 +0100
Subject: [PATCH] add test for websocket ping/pong

---
 .../frontend/ws/WebSocketEndpoint.java        | 13 ++---
 .../commafeed/integration/WebSocketIT.java    | 47 ++++++++++++++-----
 2 files changed, 37 insertions(+), 23 deletions(-)

diff --git a/commafeed-server/src/main/java/com/commafeed/frontend/ws/WebSocketEndpoint.java b/commafeed-server/src/main/java/com/commafeed/frontend/ws/WebSocketEndpoint.java
index 9bd19b94..1721ab71 100644
--- a/commafeed-server/src/main/java/com/commafeed/frontend/ws/WebSocketEndpoint.java
+++ b/commafeed-server/src/main/java/com/commafeed/frontend/ws/WebSocketEndpoint.java
@@ -8,7 +8,6 @@ import jakarta.websocket.CloseReason;
 import jakarta.websocket.CloseReason.CloseCodes;
 import jakarta.websocket.Endpoint;
 import jakarta.websocket.EndpointConfig;
-import jakarta.websocket.MessageHandler;
 import jakarta.websocket.Session;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -30,15 +29,9 @@ public class WebSocketEndpoint extends Endpoint {
 			sessions.add(userId, session);
 		}
 
-		// converting this anonymous inner class to a lambda causes the following error when a message is sent from the client
-		// Unable to find decoder for type <javax.websocket.MessageHandler$Whole>
-		// this error is only visible when registering a listener to ws.onclose on the client
-		session.addMessageHandler(new MessageHandler.Whole<String>() {
-			@Override
-			public void onMessage(String message) {
-				if ("ping".equals(message)) {
-					session.getAsyncRemote().sendText("pong");
-				}
+		session.addMessageHandler(String.class, message -> {
+			if ("ping".equals(message)) {
+				session.getAsyncRemote().sendText("pong");
 			}
 		});
 	}
diff --git a/commafeed-server/src/test/java/com/commafeed/integration/WebSocketIT.java b/commafeed-server/src/test/java/com/commafeed/integration/WebSocketIT.java
index 9e908d29..a52f5038 100644
--- a/commafeed-server/src/test/java/com/commafeed/integration/WebSocketIT.java
+++ b/commafeed-server/src/test/java/com/commafeed/integration/WebSocketIT.java
@@ -18,7 +18,6 @@ import jakarta.websocket.ContainerProvider;
 import jakarta.websocket.DeploymentException;
 import jakarta.websocket.Endpoint;
 import jakarta.websocket.EndpointConfig;
-import jakarta.websocket.MessageHandler;
 import jakarta.websocket.Session;
 
 class WebSocketIT extends BaseIT {
@@ -26,24 +25,14 @@ class WebSocketIT extends BaseIT {
 	@Test
 	void subscribeAndGetsNotified() throws DeploymentException, IOException {
 		String sessionId = login();
-		ClientEndpointConfig config = ClientEndpointConfig.Builder.create().configurator(new ClientEndpointConfig.Configurator() {
-			@Override
-			public void beforeRequest(Map<String, List<String>> headers) {
-				headers.put("Cookie", Collections.singletonList("JSESSIONID=" + sessionId));
-			}
-		}).build();
+		ClientEndpointConfig config = buildConfig(sessionId);
 
 		AtomicBoolean connected = new AtomicBoolean();
 		AtomicReference<String> messageRef = new AtomicReference<>();
 		try (Session ignored = ContainerProvider.getWebSocketContainer().connectToServer(new Endpoint() {
 			@Override
 			public void onOpen(Session session, EndpointConfig config) {
-				session.addMessageHandler(new MessageHandler.Whole<String>() {
-					@Override
-					public void onMessage(String message) {
-						messageRef.set(message);
-					}
-				});
+				session.addMessageHandler(String.class, messageRef::set);
 				connected.set(true);
 			}
 		}, config, URI.create(getWebSocketUrl()))) {
@@ -56,4 +45,36 @@ class WebSocketIT extends BaseIT {
 		}
 	}
 
+	@Test
+	void pingPong() throws DeploymentException, IOException {
+		String sessionId = login();
+		ClientEndpointConfig config = buildConfig(sessionId);
+
+		AtomicBoolean connected = new AtomicBoolean();
+		AtomicReference<String> messageRef = new AtomicReference<>();
+		try (Session session = ContainerProvider.getWebSocketContainer().connectToServer(new Endpoint() {
+			@Override
+			public void onOpen(Session session, EndpointConfig config) {
+				session.addMessageHandler(String.class, messageRef::set);
+				connected.set(true);
+			}
+		}, config, URI.create(getWebSocketUrl()))) {
+			Awaitility.await().atMost(15, TimeUnit.SECONDS).untilTrue(connected);
+
+			session.getAsyncRemote().sendText("ping");
+
+			Awaitility.await().atMost(15, TimeUnit.SECONDS).until(() -> messageRef.get() != null);
+			Assertions.assertEquals("pong", messageRef.get());
+		}
+	}
+
+	private ClientEndpointConfig buildConfig(String sessionId) {
+		return ClientEndpointConfig.Builder.create().configurator(new ClientEndpointConfig.Configurator() {
+			@Override
+			public void beforeRequest(Map<String, List<String>> headers) {
+				headers.put("Cookie", Collections.singletonList("JSESSIONID=" + sessionId));
+			}
+		}).build();
+	}
+
 }