diff --git a/application/src/test/java/org/thingsboard/server/controller/WebsocketApiTest.java b/application/src/test/java/org/thingsboard/server/controller/WebsocketApiTest.java
index aa7ffaeaf5..eab878854d 100644
--- a/application/src/test/java/org/thingsboard/server/controller/WebsocketApiTest.java
+++ b/application/src/test/java/org/thingsboard/server/controller/WebsocketApiTest.java
@@ -982,12 +982,7 @@ public class WebsocketApiTest extends AbstractControllerTest {
         List<String> keys = List.of("temperature");
         long now = System.currentTimeMillis();
 
-        // Register for 2 messages: initial entity page data + error
-        getWsClient().registerWaitForUpdate(2);
-        getWsClient().sendHistoryCmd(keys, now, TimeUnit.HOURS.toMillis(1), dtf);
-        getWsClient().waitForUpdate();
-
-        EntityDataUpdate errorUpdate = JacksonUtil.fromString(getWsClient().getLastMsg(), EntityDataUpdate.class);
+        EntityDataUpdate errorUpdate = getWsClient().sendHistoryCmd(keys, now, TimeUnit.HOURS.toMillis(1), dtf);
         assertThat(errorUpdate.getErrorCode()).isEqualTo(SubscriptionErrorCode.INTERNAL_ERROR.getCode());
         assertThat(errorUpdate.getErrorMsg()).isEqualTo(exception.getMessage());
     }
@@ -1001,12 +996,7 @@ public class WebsocketApiTest extends AbstractControllerTest {
         List<String> keys = List.of("temperature");
         long now = System.currentTimeMillis();
 
-        // Register for 2 messages: initial entity page data + error
-        getWsClient().registerWaitForUpdate(2);
-        getWsClient().subscribeTsUpdate(keys, now, TimeUnit.HOURS.toMillis(1), dtf);
-        getWsClient().waitForUpdate();
-
-        EntityDataUpdate errorUpdate = JacksonUtil.fromString(getWsClient().getLastMsg(), EntityDataUpdate.class);
+        EntityDataUpdate errorUpdate = getWsClient().subscribeTsUpdate(keys, now, TimeUnit.HOURS.toMillis(1), dtf);
         assertThat(errorUpdate.getErrorCode()).isEqualTo(SubscriptionErrorCode.INTERNAL_ERROR.getCode());
         assertThat(errorUpdate.getErrorMsg()).isEqualTo(exception.getMessage());
     }