diff --git a/application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java b/application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java index 9eac721dfb..9db4949d5f 100644 --- a/application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java +++ b/application/src/main/java/org/thingsboard/server/service/apiusage/BaseApiUsageState.java @@ -40,7 +40,8 @@ public abstract class BaseApiUsageState { private final Map gaugesReportCycles = new HashMap<>(); @Getter - private final ApiUsageState apiUsageState; + @Setter + private ApiUsageState apiUsageState; @Getter private volatile long currentCycleTs; @Getter diff --git a/application/src/main/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateService.java b/application/src/main/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateService.java index 3a351528e4..30a3af7f65 100644 --- a/application/src/main/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateService.java +++ b/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 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 stateTelemetry = new ArrayList<>(); result.forEach((apiFeature, aState) -> stateTelemetry.add(new BasicTsKvEntry(ts, new StringDataEntry(apiFeature.getApiStateKey(), aState.name())))); diff --git a/application/src/test/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateServiceTest.java b/application/src/test/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateServiceTest.java index 051214c771..a0f61a4928 100644 --- a/application/src/test/java/org/thingsboard/server/service/apiusage/DefaultTbApiUsageStateServiceTest.java +++ b/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 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 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 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 getExpectedResult() { + Map expectedResult = new HashMap<>(); + for (ApiUsageRecordKey recordKey : ApiUsageRecordKey.values()) { + if (recordKey.getApiFeature() != null) { + expectedResult.put(recordKey.getApiFeature(), ApiUsageStateValue.DISABLED); + } + } + + return expectedResult; + } } \ No newline at end of file diff --git a/common/dao-api/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateService.java b/common/dao-api/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateService.java index 3d09247c7e..9512b72ecd 100644 --- a/common/dao-api/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateService.java +++ b/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); + } diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/ApiUsageState.java b/common/data/src/main/java/org/thingsboard/server/common/data/ApiUsageState.java index f6f3d82c09..9b2a2b0046 100644 --- a/common/data/src/main/java/org/thingsboard/server/common/data/ApiUsageState.java +++ b/common/data/src/main/java/org/thingsboard/server/common/data/ApiUsageState.java @@ -86,11 +86,11 @@ public class ApiUsageState extends BaseData 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); } diff --git a/common/proto/src/main/java/org/thingsboard/server/common/util/ProtoUtils.java b/common/proto/src/main/java/org/thingsboard/server/common/util/ProtoUtils.java index 264b23f118..dd376fe746 100644 --- a/common/proto/src/main/java/org/thingsboard/server/common/util/ProtoUtils.java +++ b/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; } diff --git a/common/proto/src/main/proto/queue.proto b/common/proto/src/main/proto/queue.proto index 228a4039d2..25c1f5dd01 100644 --- a/common/proto/src/main/proto/queue.proto +++ b/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 { diff --git a/dao/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateServiceImpl.java b/dao/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateServiceImpl.java index d901ae950e..3d0e0f1309 100644 --- a/dao/src/main/java/org/thingsboard/server/dao/usagerecord/ApiUsageStateServiceImpl.java +++ b/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()); diff --git a/dao/src/test/java/org/thingsboard/server/dao/service/ApiUsageStateServiceTest.java b/dao/src/test/java/org/thingsboard/server/dao/service/ApiUsageStateServiceTest.java index 36b138f42a..d3e86a1603 100644 --- a/dao/src/test/java/org/thingsboard/server/dao/service/ApiUsageStateServiceTest.java +++ b/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); } }