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 a0f61a4928..100c98adec 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,7 +15,8 @@ */ package org.thingsboard.server.service.apiusage; -import org.jetbrains.annotations.NotNull; +import com.google.common.util.concurrent.Futures; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -25,42 +26,28 @@ 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.*; import org.thingsboard.server.common.data.id.TenantId; -import org.thingsboard.server.common.data.notification.rule.trigger.ApiUsageLimitTrigger; +import org.thingsboard.server.common.data.id.TenantProfileId; +import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileConfiguration; +import org.thingsboard.server.common.data.tenant.profile.TenantProfileData; 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.queue.discovery.PartitionService; 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; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; @ExtendWith(MockitoExtension.class) public class DefaultTbApiUsageStateServiceTest { @@ -80,122 +67,181 @@ public class DefaultTbApiUsageStateServiceTest { @Mock private InternalTelemetryService tsWsService; + @Mock + private PartitionService partitionService; + @Mock private MailExecutorService mailExecutor; - - TenantId tenantId = TenantId.fromUUID(UUID.fromString("00797a3b-7aeb-4b5b-b57a-c2a810d0f112")); @Spy @InjectMocks DefaultTbApiUsageStateService service; + private ApiUsageState dummyUsageState; + private TenantApiUsageState tenantApiUsageState; + private Tenant dummyTenant; + + private static final int MAX_ENABLE_VALUE = 5000; + private static final double WARN_THRESHOLD_VALUE = 0.8; + @BeforeEach public void setUp() { + var tenantId = TenantId.fromUUID(UUID.randomUUID()); + dummyTenant = new Tenant(); + dummyTenant.setId(tenantId); + dummyTenant.setEmail("test@tenant.com"); + + lenient().when(tsWsService.saveTimeseriesInternal(any())).thenReturn(Futures.immediateFuture(0)); + ReflectionTestUtils.setField(service, "tsWsService", tsWsService); + + DefaultTenantProfileConfiguration config = DefaultTenantProfileConfiguration.builder() + .maxJSExecutions(MAX_ENABLE_VALUE) + .maxTransportMessages(MAX_ENABLE_VALUE) + .maxRuleChains(MAX_ENABLE_VALUE) + .maxTbelExecutions(MAX_ENABLE_VALUE) + .maxDPStorageDays(MAX_ENABLE_VALUE) + .maxREExecutions(MAX_ENABLE_VALUE) + .maxEmails(MAX_ENABLE_VALUE) + .maxSms(MAX_ENABLE_VALUE) + .maxCreatedAlarms(MAX_ENABLE_VALUE) + .warnThreshold(WARN_THRESHOLD_VALUE) + .build(); + + TenantProfileData profileData = new TenantProfileData(); + profileData.setConfiguration(config); + + TenantProfile tenantProfile = new TenantProfile(); + tenantProfile.setId(new TenantProfileId(UUID.randomUUID())); + tenantProfile.setProfileData(profileData); + + dummyUsageState = new ApiUsageState(); + dummyUsageState.setTransportState(ApiUsageStateValue.ENABLED); + dummyUsageState.setDbStorageState(ApiUsageStateValue.ENABLED); + dummyUsageState.setReExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setJsExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setTbelExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setEmailExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setSmsExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setAlarmExecState(ApiUsageStateValue.ENABLED); + dummyUsageState.setTenantId(tenantId); + dummyUsageState.setEntityId(tenantId); + dummyUsageState.setVersion(1L); + + tenantApiUsageState = new TenantApiUsageState(tenantProfile, dummyUsageState); + + service.myUsageStates.put(tenantId, tenantApiUsageState); + } + + @AfterEach + public void tearDown() { + } @Test public void givenTenantIdFromEntityStatesMap_whenGetApiUsageState() { - service.myUsageStates.put(tenantId, tenantUsageStateMock); - ApiUsageState tenantUsageState = service.getApiUsageState(tenantId); + service.myUsageStates.put(dummyTenant.getTenantId(), tenantUsageStateMock); + ApiUsageState tenantUsageState = service.getApiUsageState(dummyTenant.getTenantId()); assertThat(tenantUsageState, is(tenantUsageStateMock.getApiUsageState())); - Mockito.verify(service, never()).getOrFetchState(tenantId, tenantId); + Mockito.verify(service, never()).getOrFetchState(dummyTenant.getTenantId(), dummyTenant.getTenantId()); } @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); + public void testProcess_AllFeaturesTransitionToWarning() { + TransportProtos.ToUsageStatsServiceMsg.Builder msgBuilder = TransportProtos.ToUsageStatsServiceMsg.newBuilder() + .setTenantIdMSB(dummyTenant.getId().getId().getMostSignificantBits()) + .setTenantIdLSB(dummyTenant.getId().getId().getLeastSignificantBits()) + .setCustomerIdMSB(0) + .setCustomerIdLSB(0) + .setServiceId("testService"); + + for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) { + if (key.getApiFeature() != null) { + msgBuilder.addValues(TransportProtos.UsageStatsKVProto.newBuilder() + .setKey(key.name()) + .setValue(4500L) + .build()); } } TransportProtos.ToUsageStatsServiceMsg statsMsg = msgBuilder.build(); - TbProtoQueueMsg queueMsg = new TbProtoQueueMsg<>(UUID.randomUUID(), statsMsg); + TbProtoQueueMsg msg = new TbProtoQueueMsg<>(UUID.randomUUID(), statsMsg); TbCallback callback = mock(TbCallback.class); - service.process(queueMsg, callback); - verify(callback).onSuccess(); + when(apiUsageStateService.update(any(ApiUsageState.class))).thenAnswer(invocation -> { + ApiUsageState state = invocation.getArgument(0); + state.setVersion(2L); + return state; + }); + + when(tenantService.findTenantById(dummyTenant.getTenantId())).thenReturn(dummyTenant); - 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()) - )); + service.process(msg, callback); + + verify(callback, times(1)).onSuccess(); + + for (ApiFeature feature : ApiFeature.values()) { + if (containsFeature(feature)) { + assertEquals(ApiUsageStateValue.WARNING, tenantApiUsageState.getFeatureValue(feature), + "For feature " + feature + " expected state WARNING"); + } } - } + + assertEquals(ApiUsageStateValue.WARNING, tenantApiUsageState.getFeatureValue(ApiFeature.JS)); - @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; + verify(notificationRuleProcessor, atLeastOnce()).process(any()); + verify(mailExecutor, atLeastOnce()).submit((Runnable) any()); } - private Map getExpectedResult() { - Map expectedResult = new HashMap<>(); - for (ApiUsageRecordKey recordKey : ApiUsageRecordKey.values()) { - if (recordKey.getApiFeature() != null) { - expectedResult.put(recordKey.getApiFeature(), ApiUsageStateValue.DISABLED); + @Test + public void testProcess_AllFeaturesTransitionToDisabled() { + TransportProtos.ToUsageStatsServiceMsg.Builder msgBuilder = TransportProtos.ToUsageStatsServiceMsg.newBuilder() + .setTenantIdMSB(dummyTenant.getId().getId().getMostSignificantBits()) + .setTenantIdLSB(dummyTenant.getId().getId().getLeastSignificantBits()) + .setCustomerIdMSB(0) + .setCustomerIdLSB(0) + .setServiceId("testService"); + + for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) { + if (key.getApiFeature() != null) { + msgBuilder.addValues(TransportProtos.UsageStatsKVProto.newBuilder() + .setKey(key.name()) + .setValue(5500L) + .build()); + } + } + TransportProtos.ToUsageStatsServiceMsg statsMsg = msgBuilder.build(); + TbProtoQueueMsg msg = new TbProtoQueueMsg<>(UUID.randomUUID(), statsMsg); + TbCallback callback = mock(TbCallback.class); + + when(apiUsageStateService.update(any(ApiUsageState.class))).thenAnswer(invocation -> { + ApiUsageState state = invocation.getArgument(0); + state.setVersion(2L); + return state; + }); + when(tenantService.findTenantById(dummyTenant.getTenantId())).thenReturn(dummyTenant); + + service.process(msg, callback); + verify(callback, times(1)).onSuccess(); + + for (ApiFeature feature : ApiFeature.values()) { + if (containsFeature(feature)) { + assertEquals(ApiUsageStateValue.DISABLED, tenantApiUsageState.getFeatureValue(feature), + "For feature " + feature + " expected state DISABLED"); } } + verify(notificationRuleProcessor, atLeastOnce()).process(any()); + verify(mailExecutor, atLeastOnce()).submit((Runnable) any()); + } - return expectedResult; + private boolean containsFeature(ApiFeature feature) { + for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) { + if (key.getApiFeature() != null && key.getApiFeature().equals(feature)) { + return true; + } + } + return false; } + } \ No newline at end of file