From 9b51bf7b997e71f676963e91559512c835022a70 Mon Sep 17 00:00:00 2001
From: YevhenBondarenko <ybondarenko@thingsboard.io>
Date: Thu, 30 Nov 2023 16:42:31 +0100
Subject: [PATCH] fixed flaky test

---
 .../controller/DeviceControllerTest.java      | 44 ++++++++++---------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/application/src/test/java/org/thingsboard/server/controller/DeviceControllerTest.java b/application/src/test/java/org/thingsboard/server/controller/DeviceControllerTest.java
index cefe24e935..7cafe2bbcb 100644
--- a/application/src/test/java/org/thingsboard/server/controller/DeviceControllerTest.java
+++ b/application/src/test/java/org/thingsboard/server/controller/DeviceControllerTest.java
@@ -28,12 +28,14 @@ import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.mockito.AdditionalAnswers;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.mock.mockito.SpyBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Primary;
 import org.springframework.test.context.ContextConfiguration;
+import org.testcontainers.shaded.org.awaitility.Awaitility;
 import org.thingsboard.common.util.JacksonUtil;
 import org.thingsboard.common.util.ThingsBoardExecutors;
 import org.thingsboard.server.common.data.Customer;
@@ -72,6 +74,7 @@ import org.thingsboard.server.dao.exception.DataValidationException;
 import org.thingsboard.server.dao.exception.DeviceCredentialsValidationException;
 import org.thingsboard.server.dao.model.ModelConstants;
 import org.thingsboard.server.dao.service.DaoSqlTest;
+import org.thingsboard.server.gen.transport.TransportProtos;
 import org.thingsboard.server.service.gateway_device.GatewayNotificationsService;
 import org.thingsboard.server.service.state.DeviceStateService;
 
@@ -83,7 +86,6 @@ import java.util.concurrent.TimeUnit;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
@@ -1345,27 +1347,27 @@ public class DeviceControllerTest extends AbstractControllerTest {
                 ActionType.ASSIGNED_TO_TENANT, savedDifferentTenant.getId().getId().toString(), savedDifferentTenant.getTitle());
         testNotificationUpdateGatewayNever();
 
-        Mockito.verify(deviceStateService, times(1)).onQueueMsg(
-                argThat(proto ->
-                        proto.getTenantIdMSB() == savedTenant.getUuidId().getMostSignificantBits() &&
-                                proto.getTenantIdLSB() == savedTenant.getUuidId().getLeastSignificantBits() &&
-                                proto.getDeviceIdMSB() == savedDevice.getUuidId().getMostSignificantBits() &&
-                                proto.getDeviceIdLSB() == savedDevice.getUuidId().getLeastSignificantBits() &&
-                                proto.getDeleted()
-                ),
-                any()
-        );
+        ArgumentCaptor<TransportProtos.DeviceStateServiceMsgProto> protoCaptor = ArgumentCaptor.forClass(TransportProtos.DeviceStateServiceMsgProto.class);
 
-        Mockito.verify(deviceStateService, times(1)).onQueueMsg(
-                argThat(proto ->
-                        proto.getTenantIdMSB() == savedDifferentTenant.getUuidId().getMostSignificantBits() &&
-                                proto.getTenantIdLSB() == savedDifferentTenant.getUuidId().getLeastSignificantBits() &&
-                                proto.getDeviceIdMSB() == savedDevice.getUuidId().getMostSignificantBits() &&
-                                proto.getDeviceIdLSB() == savedDevice.getUuidId().getLeastSignificantBits() &&
-                                proto.getAdded()
-                ),
-                any()
-        );
+        Awaitility.await().atMost(5, TimeUnit.SECONDS).until(() -> {
+            Mockito.verify(deviceStateService, Mockito.atLeastOnce()).onQueueMsg(protoCaptor.capture(), any());
+            return protoCaptor.getAllValues().stream().anyMatch(proto ->
+                    proto.getTenantIdMSB() == savedTenant.getUuidId().getMostSignificantBits() &&
+                            proto.getTenantIdLSB() == savedTenant.getUuidId().getLeastSignificantBits() &&
+                            proto.getDeviceIdMSB() == savedDevice.getUuidId().getMostSignificantBits() &&
+                            proto.getDeviceIdLSB() == savedDevice.getUuidId().getLeastSignificantBits() &&
+                            proto.getDeleted());
+        });
+
+        Awaitility.await().atMost(5, TimeUnit.SECONDS).until(() -> {
+            Mockito.verify(deviceStateService, Mockito.atLeastOnce()).onQueueMsg(protoCaptor.capture(), any());
+            return protoCaptor.getAllValues().stream().anyMatch(proto ->
+                    proto.getTenantIdMSB() == savedDifferentTenant.getUuidId().getMostSignificantBits() &&
+                            proto.getTenantIdLSB() == savedDifferentTenant.getUuidId().getLeastSignificantBits() &&
+                            proto.getDeviceIdMSB() == savedDevice.getUuidId().getMostSignificantBits() &&
+                            proto.getDeviceIdLSB() == savedDevice.getUuidId().getLeastSignificantBits() &&
+                            proto.getAdded());
+        });
 
         login("tenant9@thingsboard.org", "testPassword1");