diff --git a/application/src/test/java/org/thingsboard/server/edge/BaseAlarmEdgeTest.java b/application/src/test/java/org/thingsboard/server/edge/BaseAlarmEdgeTest.java
index 49784c3939..d952e478a0 100644
--- a/application/src/test/java/org/thingsboard/server/edge/BaseAlarmEdgeTest.java
+++ b/application/src/test/java/org/thingsboard/server/edge/BaseAlarmEdgeTest.java
@@ -73,13 +73,13 @@ abstract public class BaseAlarmEdgeTest extends AbstractEdgeTest {
     @Test
     public void testAlarms() throws Exception {
         // create alarm
-        edgeImitator.expectMessageAmount(1);
         Device device = findDeviceByName("Edge Device 1");
         Alarm alarm = new Alarm();
         alarm.setOriginator(device.getId());
         alarm.setStatus(AlarmStatus.ACTIVE_UNACK);
         alarm.setType("alarm");
         alarm.setSeverity(AlarmSeverity.CRITICAL);
+        edgeImitator.expectMessageAmount(1);
         Alarm savedAlarm = doPost("/api/alarm", alarm, Alarm.class);
         Assert.assertTrue(edgeImitator.waitForMessages());
         AbstractMessage latestMessage = edgeImitator.getLatestMessage();
diff --git a/application/src/test/java/org/thingsboard/server/edge/BaseRelationEdgeTest.java b/application/src/test/java/org/thingsboard/server/edge/BaseRelationEdgeTest.java
index d26680a01d..6a6474e5a3 100644
--- a/application/src/test/java/org/thingsboard/server/edge/BaseRelationEdgeTest.java
+++ b/application/src/test/java/org/thingsboard/server/edge/BaseRelationEdgeTest.java
@@ -39,7 +39,6 @@ abstract public class BaseRelationEdgeTest extends AbstractEdgeTest {
     @Test
     public void testRelations() throws Exception {
         // create relation
-        edgeImitator.expectMessageAmount(1);
         Device device = findDeviceByName("Edge Device 1");
         Asset asset = findAssetByName("Edge Asset 1");
         EntityRelation relation = new EntityRelation();
@@ -47,6 +46,7 @@ abstract public class BaseRelationEdgeTest extends AbstractEdgeTest {
         relation.setFrom(device.getId());
         relation.setTo(asset.getId());
         relation.setTypeGroup(RelationTypeGroup.COMMON);
+        edgeImitator.expectMessageAmount(1);
         doPost("/api/relation", relation);
         Assert.assertTrue(edgeImitator.waitForMessages());
         AbstractMessage latestMessage = edgeImitator.getLatestMessage();