Improve tests

application/src/test/java/org/thingsboard/server/controller/ApiKeyWebSocketApiTest.java

@@ -17,13 +17,23 @@ package org.thingsboard.server.controller;
 
 import lombok.extern.slf4j.Slf4j;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.thingsboard.server.common.data.pat.ApiKey;
 import org.thingsboard.server.common.data.pat.ApiKeyInfo;
+import org.thingsboard.server.common.data.query.DeviceTypeFilter;
+import org.thingsboard.server.common.data.query.EntityCountQuery;
 import org.thingsboard.server.dao.service.DaoSqlTest;
+import org.thingsboard.server.service.ws.telemetry.cmd.v2.EntityCountCmd;
+import org.thingsboard.server.service.ws.telemetry.cmd.v2.EntityCountUpdate;
 
+import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
@@ -59,6 +69,51 @@ public class ApiKeyWebSocketApiTest extends WebSocketApiTest {
         TbTestWebSocketClient client = buildAndConnectWebSocketClientWithApiKeyHeader(apiKey.getValue());
         try {
             assertThat(client.isOpen()).isTrue();
+
+            DeviceTypeFilter dtf = new DeviceTypeFilter(List.of("default"), "Device");
+            EntityCountQuery ecq = new EntityCountQuery(dtf, Collections.emptyList());
+            EntityCountCmd cmd = new EntityCountCmd(1, ecq);
+            client.send(cmd);
+
+            EntityCountUpdate update = client.parseCountReply(client.waitForReply());
+            Assert.assertEquals(1, update.getCmdId());
+            Assert.assertTrue(update.getCount() >= 0);
+        } finally {
+            client.close();
+        }
+    }
+
+    @Test
+    public void testInvalidApiKeyHeader_connectionRejected() throws Exception {
+        TbTestWebSocketClient client = new TbTestWebSocketClient(
+                new URI(WS_URL + wsPort + "/api/ws"), Map.of("X-API-Key", "invalid-key"));
+        try {
+            boolean connected = client.connectBlocking(TIMEOUT, TimeUnit.SECONDS);
+            assertThat(connected).isFalse();
+        } finally {
+            client.close();
+        }
+    }
+
+    @Test
+    public void testEmptyApiKeyHeader_connectionRejected() throws Exception {
+        TbTestWebSocketClient client = new TbTestWebSocketClient(
+                new URI(WS_URL + wsPort + "/api/ws"), Map.of("X-API-Key", ""));
+        try {
+            boolean connected = client.connectBlocking(TIMEOUT, TimeUnit.SECONDS);
+            assertThat(connected).isFalse();
+        } finally {
+            client.close();
+        }
+    }
+
+    @Test
+    public void testInvalidApiKeyAuthCmd_connectionClosed() throws Exception {
+        TbTestWebSocketClient client = new TbTestWebSocketClient(new URI(WS_URL + wsPort + "/api/ws"));
+        assertThat(client.connectBlocking(TIMEOUT, TimeUnit.SECONDS)).isTrue();
+        try {
+            client.authenticateWithApiKey("invalid-key");
+            assertThat(client.waitForClose()).isTrue();
         } finally {
             client.close();
         }

application/src/test/java/org/thingsboard/server/controller/TbTestWebSocketClient.java

@@ -53,6 +53,8 @@ public class TbTestWebSocketClient extends WebSocketClient {
 
     private static final long TIMEOUT = TimeUnit.SECONDS.toMillis(30);
 
+    private final CountDownLatch closeLatch = new CountDownLatch(1);
+
     @Getter
     private volatile String lastMsg;
     private volatile CountDownLatch reply;
@@ -98,6 +100,16 @@ public class TbTestWebSocketClient extends WebSocketClient {
     @Override
     public void onClose(int i, String s, boolean b) {
         log.info("CLOSED.");
+        closeLatch.countDown();
+    }
+
+    public boolean waitForClose() {
+        try {
+            return closeLatch.await(TIMEOUT, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            log.warn("Failed to await close", e);
+            return false;
+        }
     }
 
     @Override