fixed flaky edqs tests

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

@@ -25,6 +25,9 @@ import org.thingsboard.server.common.data.edqs.EdqsState;
 import org.thingsboard.server.common.data.edqs.EdqsState.EdqsApiMode;
 import org.thingsboard.server.common.data.edqs.ToCoreEdqsRequest;
 import org.thingsboard.server.common.data.page.PageData;
+import org.thingsboard.server.common.data.query.AlarmCountQuery;
+import org.thingsboard.server.common.data.query.AlarmData;
+import org.thingsboard.server.common.data.query.AlarmDataQuery;
 import org.thingsboard.server.common.data.query.EntityCountQuery;
 import org.thingsboard.server.common.data.query.EntityData;
 import org.thingsboard.server.common.data.query.EntityDataQuery;
@@ -70,12 +73,24 @@ public class EdqsEntityQueryControllerTest extends EntityQueryControllerTest {
                 result -> result.getTotalElements() == expectedResultSize);
     }
 
+    @Override
+    protected PageData<AlarmData> findAlarmsByQueryAndCheck(AlarmDataQuery query, int expectedResultSize) {
+        return await().atMost(TIMEOUT, TimeUnit.SECONDS).until(() -> findAlarmsByQuery(query),
+                result -> result.getTotalElements() == expectedResultSize);
+    }
+
     @Override
     protected Long countByQueryAndCheck(EntityCountQuery query, long expectedResult) {
         return await().atMost(TIMEOUT, TimeUnit.SECONDS).until(() -> countByQuery(query),
                 result -> result == expectedResult);
     }
 
+    @Override
+    protected Long countAlarmsByQueryAndCheck(AlarmCountQuery query, long expectedResult) {
+        return await().atMost(TIMEOUT, TimeUnit.SECONDS).until(() -> countAlarmsByQuery(query),
+                result -> result == expectedResult);
+    }
+
     @Test
     public void testEdqsState() throws Exception {
         loginSysAdmin();