Resolved review comments

application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java

@@ -40,7 +40,8 @@ public abstract class BaseApiUsageState {
     private final Map<ApiUsageRecordKey, Long> gaugesReportCycles = new HashMap<>();
 
     @Getter
-    private final ApiUsageState apiUsageState;
+    @Setter
+    private ApiUsageState apiUsageState;
     @Getter
     private volatile long currentCycleTs;
     @Getter

application/src/main/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateService.java

@@ -355,7 +355,8 @@ public class DefaultTbApiUsageStateService extends AbstractPartitionBasedService
 
     private void persistAndNotify(BaseApiUsageState state, Map<ApiFeature, ApiUsageStateValue> result) {
         log.info("[{}] Detected update of the API state for {}: {}", state.getEntityId(), state.getEntityType(), result);
-        apiUsageStateService.update(state.getApiUsageState());
+        ApiUsageState updatedState = apiUsageStateService.update(state.getApiUsageState());
+        state.setApiUsageState(updatedState);
         long ts = System.currentTimeMillis();
         List<TsKvEntry> stateTelemetry = new ArrayList<>();
         result.forEach((apiFeature, aState) -> stateTelemetry.add(new BasicTsKvEntry(ts, new StringDataEntry(apiFeature.getApiStateKey(), aState.name()))));

application/src/test/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateServiceTest.java

@@ -15,6 +15,7 @@
  */
 package org.thingsboard.server.service.apiusage;
 
+import org.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -23,14 +24,43 @@ import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.Spy;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.thingsboard.server.common.data.ApiFeature;
+import org.thingsboard.server.common.data.ApiUsageRecordKey;
 import org.thingsboard.server.common.data.ApiUsageState;
+import org.thingsboard.server.common.data.ApiUsageStateValue;
+import org.thingsboard.server.common.data.EntityType;
+import org.thingsboard.server.common.data.Tenant;
 import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.data.notification.rule.trigger.ApiUsageLimitTrigger;
+import org.thingsboard.server.common.msg.notification.NotificationRuleProcessor;
+import org.thingsboard.server.common.msg.queue.TbCallback;
+import org.thingsboard.server.dao.tenant.TenantService;
+import org.thingsboard.server.dao.usagerecord.ApiUsageStateService;
+import org.thingsboard.server.gen.transport.TransportProtos;
+import org.thingsboard.server.queue.common.TbProtoQueueMsg;
+import org.thingsboard.server.service.mail.MailExecutorService;
+import org.thingsboard.server.service.telemetry.InternalTelemetryService;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
 public class DefaultTbApiUsageStateServiceTest {
@@ -38,6 +68,21 @@ public class DefaultTbApiUsageStateServiceTest {
     @Mock
     TenantApiUsageState tenantUsageStateMock;
 
+    @Mock
+    private NotificationRuleProcessor notificationRuleProcessor;
+
+    @Mock
+    private ApiUsageStateService apiUsageStateService;
+
+    @Mock
+    private TenantService tenantService;
+
+    @Mock
+    private InternalTelemetryService tsWsService;
+
+    @Mock
+    private MailExecutorService mailExecutor;
+    
     TenantId tenantId = TenantId.fromUUID(UUID.fromString("00797a3b-7aeb-4b5b-b57a-c2a810d0f112"));
 
     @Spy
@@ -46,6 +91,7 @@ public class DefaultTbApiUsageStateServiceTest {
 
     @BeforeEach
     public void setUp() {
+        ReflectionTestUtils.setField(service, "tsWsService", tsWsService);
     }
 
     @Test
@@ -56,4 +102,100 @@ public class DefaultTbApiUsageStateServiceTest {
         Mockito.verify(service, never()).getOrFetchState(tenantId, tenantId);
     }
 
+    @Test
+    public void testAllApiFeaturesDisabledWhenLimitReached() {
+        doReturn(null).when(tsWsService).saveTimeseriesInternal(any());
+
+        TenantApiUsageState tenantUsageStateMock = mock(TenantApiUsageState.class);
+        ApiUsageState apiUsageState = getApiUsageState();
+        when(tenantUsageStateMock.getApiUsageState()).thenReturn(apiUsageState);
+
+        doReturn(BaseApiUsageState.StatsCalculationResult.builder()
+                .newValue(200L)
+                .valueChanged(true)
+                .newHourlyValue(200L)
+                .hourlyValueChanged(true)
+                .build())
+                .when(tenantUsageStateMock).calculate(any(ApiUsageRecordKey.class), anyLong(), anyString());
+
+        doReturn(200L).when(tenantUsageStateMock).getProfileThreshold(any(ApiUsageRecordKey.class));
+        doReturn(50L).when(tenantUsageStateMock).getProfileWarnThreshold(any(ApiUsageRecordKey.class));
+        doReturn(300L).when(tenantUsageStateMock).get(any(ApiUsageRecordKey.class));
+
+        when(tenantUsageStateMock.getEntityType()).thenReturn(EntityType.TENANT);
+        when(tenantUsageStateMock.getEntityId()).thenReturn(tenantId);
+
+        Map<ApiFeature, ApiUsageStateValue> expectedResult = getExpectedResult();
+        when(tenantUsageStateMock.checkStateUpdatedDueToThreshold(any())).thenReturn(expectedResult);
+        service.myUsageStates.put(tenantId, tenantUsageStateMock);
+
+        when(apiUsageStateService.update(apiUsageState)).thenReturn(apiUsageState);
+
+        Tenant dummyTenant = new Tenant();
+        dummyTenant.setEmail("test@example.com");
+        when(tenantService.findTenantById(any())).thenReturn(dummyTenant);
+
+        TransportProtos.ToUsageStatsServiceMsg.Builder msgBuilder = TransportProtos.ToUsageStatsServiceMsg.newBuilder();
+        UUID uuid = tenantId.getId();
+        msgBuilder.setTenantIdMSB(uuid.getMostSignificantBits())
+                .setTenantIdLSB(uuid.getLeastSignificantBits());
+        msgBuilder.setCustomerIdMSB(0)
+                .setCustomerIdLSB(0);
+        msgBuilder.setServiceId("TEST_SERVICE");
+
+        List<TransportProtos.UsageStatsKVProto> usageStats = new ArrayList<>();
+        for (ApiUsageRecordKey recordKey : ApiUsageRecordKey.values()) {
+            if (recordKey.getApiFeature() != null) {
+                TransportProtos.UsageStatsKVProto stat = TransportProtos.UsageStatsKVProto.newBuilder()
+                        .setKey(recordKey.name())
+                        .setValue(1000L)
+                        .build();
+                usageStats.add(stat);
+                msgBuilder.addValues(stat);
+            }
+        }
+
+        TransportProtos.ToUsageStatsServiceMsg statsMsg = msgBuilder.build();
+        TbProtoQueueMsg<TransportProtos.ToUsageStatsServiceMsg> queueMsg = new TbProtoQueueMsg<>(UUID.randomUUID(), statsMsg);
+        TbCallback callback = mock(TbCallback.class);
+
+        service.process(queueMsg, callback);
+        verify(callback).onSuccess();
+
+        for (ApiFeature feature : expectedResult.keySet()) {
+            verify(notificationRuleProcessor, atLeastOnce()).process(argThat(trigger ->
+                    trigger instanceof ApiUsageLimitTrigger &&
+                            ((ApiUsageLimitTrigger) trigger).getStatus() == ApiUsageStateValue.DISABLED &&
+                            ((ApiUsageLimitTrigger) trigger).getState().getApiFeature().getApiStateKey().equals(feature.getApiStateKey())
+            ));
+        }
+    }
+
+
+    @NotNull
+    private ApiUsageState getApiUsageState() {
+        ApiUsageState apiUsageState = new ApiUsageState();
+
+        apiUsageState.setTenantId(tenantId);
+        apiUsageState.setTransportState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setDbStorageState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setReExecState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setJsExecState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setTbelExecState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setEmailExecState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setSmsExecState(ApiUsageStateValue.ENABLED);
+        apiUsageState.setAlarmExecState(ApiUsageStateValue.ENABLED);
+        return apiUsageState;
+    }
+
+    private Map<ApiFeature, ApiUsageStateValue> getExpectedResult() {
+        Map<ApiFeature, ApiUsageStateValue> expectedResult = new HashMap<>();
+        for (ApiUsageRecordKey recordKey : ApiUsageRecordKey.values()) {
+            if (recordKey.getApiFeature() != null) {
+                expectedResult.put(recordKey.getApiFeature(), ApiUsageStateValue.DISABLED);
+            }
+        }
+
+        return expectedResult;
+    }
 }

common/dao-api/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateService.java

@@ -34,4 +34,5 @@ public interface ApiUsageStateService extends EntityDaoService {
     void deleteApiUsageStateByEntityId(EntityId entityId);
 
     ApiUsageState findApiUsageStateById(TenantId tenantId, ApiUsageStateId id);
+
 }

common/data/src/main/java/org/thingsboard/server/common/data/ApiUsageState.java

@@ -86,11 +86,11 @@ public class ApiUsageState extends BaseData<ApiUsageStateId> implements HasTenan
         return !ApiUsageStateValue.DISABLED.equals(tbelExecState);
     }
 
-    public boolean isEmailSendEnabled(){
+    public boolean isEmailSendEnabled() {
         return !ApiUsageStateValue.DISABLED.equals(emailExecState);
     }
 
-    public boolean isSmsSendEnabled(){
+    public boolean isSmsSendEnabled() {
         return !ApiUsageStateValue.DISABLED.equals(smsExecState);
     }
 

common/proto/src/main/java/org/thingsboard/server/common/util/ProtoUtils.java

@@ -1019,7 +1019,9 @@ public class ProtoUtils {
                 .setTbelExecState(apiUsageState.getTbelExecState().name())
                 .setEmailExecState(apiUsageState.getEmailExecState().name())
                 .setSmsExecState(apiUsageState.getSmsExecState().name())
-                .setAlarmExecState(apiUsageState.getAlarmExecState().name()).build();
+                .setAlarmExecState(apiUsageState.getAlarmExecState().name())
+                .setVersion(apiUsageState.getVersion())
+                .build();
     }
 
     public static ApiUsageState fromProto(TransportProtos.ApiUsageStateProto proto) {
@@ -1035,6 +1037,7 @@ public class ProtoUtils {
         apiUsageState.setEmailExecState(ApiUsageStateValue.valueOf(proto.getEmailExecState()));
         apiUsageState.setSmsExecState(ApiUsageStateValue.valueOf(proto.getSmsExecState()));
         apiUsageState.setAlarmExecState(ApiUsageStateValue.valueOf(proto.getAlarmExecState()));
+        apiUsageState.setVersion(proto.getVersion());
         return apiUsageState;
     }
 

common/proto/src/main/proto/queue.proto

@@ -325,6 +325,7 @@ message ApiUsageStateProto {
   string emailExecState = 14;
   string smsExecState = 15;
   string alarmExecState = 16;
+  int64 version = 17;
 }
 
 message RepositorySettingsProto {

dao/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateServiceImpl.java

@@ -154,6 +154,7 @@ public class ApiUsageStateServiceImpl extends AbstractEntityService implements A
         log.trace("Executing save [{}]", apiUsageState.getTenantId());
         validateId(apiUsageState.getTenantId(), id -> INCORRECT_TENANT_ID + id);
         validateId(apiUsageState.getId(), "Can't save new usage state. Only update is allowed!");
+        apiUsageState.setVersion(null);
         ApiUsageState savedState = apiUsageStateDao.save(apiUsageState.getTenantId(), apiUsageState);
         eventPublisher.publishEvent(SaveEntityEvent.builder().tenantId(savedState.getTenantId()).entityId(savedState.getId())
                 .entity(savedState).build());

dao/src/test/java/org/thingsboard/server/dao/service/ApiUsageStateServiceTest.java

@@ -20,6 +20,7 @@ import org.junit.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.thingsboard.server.common.data.ApiUsageState;
 import org.thingsboard.server.common.data.ApiUsageStateValue;
+import org.thingsboard.server.dao.exception.IncorrectParameterException;
 import org.thingsboard.server.dao.usagerecord.ApiUsageStateService;
 
 
@@ -30,22 +31,42 @@ public class ApiUsageStateServiceTest extends AbstractServiceTest {
     ApiUsageStateService apiUsageStateService;
 
     @Test
-    public void testFindApiUsageStateByTenantId() {
-        ApiUsageState apiUsageState = apiUsageStateService.findTenantApiUsageState(tenantId);
-        Assert.assertNotNull(apiUsageState);
+    public void testFindTenantApiUsageState() {
+        ApiUsageState state = apiUsageStateService.findTenantApiUsageState(tenantId);
+        Assert.assertNotNull(state);
     }
 
     @Test
-    public void testUpdateApiUsageState(){
-        ApiUsageState apiUsageState = apiUsageStateService.findTenantApiUsageState(tenantId);
-        Assert.assertNotNull(apiUsageState);
-        Assert.assertTrue(apiUsageState.isTransportEnabled());
-        apiUsageState.setTransportState(ApiUsageStateValue.DISABLED);
-        apiUsageState = apiUsageStateService.update(apiUsageState);
-        Assert.assertNotNull(apiUsageState);
-        apiUsageState = apiUsageStateService.findTenantApiUsageState(tenantId);
-        Assert.assertNotNull(apiUsageState);
-        Assert.assertFalse(apiUsageState.isTransportEnabled());
+    public void testUpdate() {
+        ApiUsageState state = apiUsageStateService.findTenantApiUsageState(tenantId);
+
+        state.setTransportState(ApiUsageStateValue.DISABLED);
+        ApiUsageState updated = apiUsageStateService.update(state);
+        Assert.assertEquals(ApiUsageStateValue.DISABLED, updated.getTransportState());
+    }
+
+    @Test
+    public void testUpdateWithNullId() {
+        ApiUsageState newState = new ApiUsageState();
+        newState.setTenantId(tenantId);
+        newState.setTransportState(ApiUsageStateValue.ENABLED);
+        Assert.assertThrows(IncorrectParameterException.class, () -> apiUsageStateService.update(newState));
+    }
+
+    @Test
+    public void testFindApiUsageStateByEntityId() {
+        ApiUsageState state = apiUsageStateService.findApiUsageStateByEntityId(tenantId);
+        Assert.assertNotNull(state);
+    }
+
+    @Test
+    public void testDeleteByTenantId() {
+        ApiUsageState state = apiUsageStateService.findTenantApiUsageState(tenantId);
+        Assert.assertNotNull(state);
+
+        apiUsageStateService.deleteByTenantId(tenantId);
+        state = apiUsageStateService.findTenantApiUsageState(tenantId);
+        Assert.assertNull(state);
     }
 
 }