fix tests

Author: guFalcon

src/main/java/info/unterrainer/websocketserver/WebsocketServer.java

@@ -155,7 +155,7 @@ public WebsocketServer wsOauth(String path, WsOauthHandlerBase handler) {
 		handler.setTokenHandler(tokenManager);
 		wss.ws(path, ws -> {
 			ws.onConnect(handler::onConnect);
-			ws.onMessage(handler::onMessage);
+			ws.onMessage(handler::onMsg);
 			ws.onBinaryMessage(handler::onBinaryMsg);
 			ws.onClose(handler::onClose);
 			ws.onError(handler::onError);

src/main/java/info/unterrainer/websocketserver/WsOauthHandlerBase.java

@@ -69,6 +69,7 @@ public void removeClient(Session session) {
 		log.debug("(" + name + ") Removing client: [{}]", session.getRemoteAddress());
 		clientsConnected.removeIf(client -> client.session.equals(session));
 		clientsQuarantined.removeIf(client -> client.session.equals(session));
+		tenantIdsBySession.remove(session);
 	}
 
 	public WsConnectContext getClient(Session session) {
@@ -93,6 +94,10 @@ public boolean isConnected(Session session) {
 
 	@Override
 	public void onConnect(WsConnectContext ctx) throws Exception {
+		handleConnect(ctx);
+	}
+	
+	protected void handleConnect(WsConnectContext ctx) throws Exception {
 		log.debug("(" + name + ") New client tries to connect: [{}]", ctx.session.getRemoteAddress());
 		String token = ctx.header("Authorization");
 		if (token == null || token.isEmpty()) {
@@ -116,9 +121,10 @@ public void onConnect(WsConnectContext ctx) throws Exception {
 
 	@Override
 	public void onMsg(WsMessageContext ctx) throws Exception {
+		handleMessage(ctx);
 	}
 
-	public final void onMessage(WsMessageContext ctx) throws Exception {
+	protected void handleMessage(WsMessageContext ctx) throws Exception {
 		log.debug("(" + name + ") Received from [{}]: [{}]", ctx.session.getRemoteAddress(), ctx.message());
 		if (isQuarantined(ctx.session)) {
 			log.warn(
@@ -153,9 +159,10 @@ public final void onMessage(WsMessageContext ctx) throws Exception {
 
 	@Override
 	public void onBinaryMsg(WsBinaryMessageContext ctx) throws Exception {
+		handleBinaryMessage(ctx);
 	}
 
-	public void onBinaryMessage(WsBinaryMessageContext ctx) throws Exception {
+	protected void handleBinaryMessage(WsBinaryMessageContext ctx) throws Exception {
 		log.debug("(" + name + ") Received binary message from [{}]: [{}] bytes", ctx.session.getRemoteAddress(),
 				ctx.data().length);
 		if (isQuarantined(ctx.session)) {
@@ -170,12 +177,20 @@ public void onBinaryMessage(WsBinaryMessageContext ctx) throws Exception {
 
 	@Override
 	public void onClose(WsCloseContext ctx) throws Exception {
+		handleClose(ctx);
+	}
+	
+	protected void handleClose(WsCloseContext ctx) throws Exception {
 		log.debug("(" + name + ") Disconnected client: [{}]", ctx.session.getRemoteAddress());
 		removeClient(ctx.session);
 	}
 
 	@Override
 	public void onError(WsErrorContext ctx) throws Exception {
+		handleError(ctx);
+	}
+	
+	protected void handleError(WsErrorContext ctx) throws Exception {
 		Throwable t = ctx.error();
 		if (t instanceof EOFException || t instanceof IOException) {
 			log.debug("(" + name + ") Client disconnected [{}].", ctx.session.getRemoteAddress());

src/test/java/info/unterrainer/websocketserver/WebsocketReconnectionIntegrationTest.java

@@ -97,6 +97,7 @@ void testClientDisconnectsAndReconnects() throws Exception {
 		assertThat(connectionIds).hasSize(2);
 
 		session2.close();
+		server.stop();
 	}
 
 	@Test
@@ -145,14 +146,14 @@ void testServerForcesDisconnectAndClientReconnects() throws Exception {
 		assertThat(session2.isOpen()).isTrue();
 
 		session2.close();
+		server.stop();
 	}
 
 	@Test
 	void testMessageLossOnReconnection() throws Exception {
 		server = new WebsocketServer("message-loss-server");
 		CountDownLatch firstConnect = new CountDownLatch(1);
 		CountDownLatch secondConnect = new CountDownLatch(1);
-		AtomicInteger messageCounter = new AtomicInteger(0);
 		List<Integer> receivedMessages = new ArrayList<>();
 
 		server.ws("/message-loss", ws -> {
@@ -211,6 +212,7 @@ void testMessageLossOnReconnection() throws Exception {
 		assertThat(receivedMessages).containsExactly(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
 
 		session2.close();
+		server.stop();
 	}
 
 	@Test
@@ -250,6 +252,7 @@ void testRapidReconnectionAttempts() throws Exception {
 
 		assertThat(connectionCount.get()).isEqualTo(attempts);
 		assertThat(disconnectionCount.get()).isEqualTo(attempts);
+		server.stop();
 	}
 
 	@Test
@@ -297,6 +300,7 @@ void testReconnectionWithDifferentClientInstance() throws Exception {
 
 		session2.close();
 		client2.stop();
+		server.stop();
 	}
 
 	@Test
@@ -366,6 +370,7 @@ void testReconnectionPreservesSessionState() throws Exception {
 		assertThat(firstSessionId.get()).isNotEqualTo(secondSessionId.get());
 
 		session2.close();
+		server.stop();
 	}
 
 	@Test
@@ -399,6 +404,7 @@ void testNetworkInterruptionSimulation() throws Exception {
 
 		// Handler should detect the closure
 		assertThat(handler.closeLatch.await(5, TimeUnit.SECONDS)).isTrue();
+		server.stop();
 	}
 
 	@WebSocket

src/test/java/info/unterrainer/websocketserver/WsHandlerBaseTest.java

@@ -31,6 +31,12 @@ class WsHandlerBaseTest {
 
 	@BeforeEach
 	void setUp() {
+		// Initialize mocks explicitly to avoid relying on Mockito annotations processing
+		connectContext = mock(WsConnectContext.class);
+		messageContext = mock(WsMessageContext.class);
+		binaryMessageContext = mock(WsBinaryMessageContext.class);
+		closeContext = mock(WsCloseContext.class);
+		errorContext = mock(WsErrorContext.class);
 		handler = new TestWsHandler();
 	}
 
@@ -157,4 +163,4 @@ public void onError(WsErrorContext ctx) throws Exception {
 			lastError = ctx.error();
 		}
 	}
-}
+}