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 c9abb8e1c0..8fefde0f6c 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 @@ -16,13 +16,12 @@ package org.thingsboard.server.service.apiusage; import com.google.common.util.concurrent.Futures; -import org.junit.jupiter.api.AfterEach; +import org.jetbrains.annotations.NotNull; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InjectMocks; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.Spy; import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.test.util.ReflectionTestUtils; @@ -33,6 +32,8 @@ import org.thingsboard.server.common.data.tenant.profile.DefaultTenantProfileCon 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.common.msg.queue.TopicPartitionInfo; +import org.thingsboard.server.dao.tenant.TbTenantProfileCache; import org.thingsboard.server.dao.tenant.TenantService; import org.thingsboard.server.dao.usagerecord.ApiUsageStateService; import org.thingsboard.server.gen.transport.TransportProtos; @@ -46,8 +47,11 @@ import java.util.UUID; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -58,9 +62,6 @@ import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) public class DefaultTbApiUsageStateServiceTest { - @Mock - TenantApiUsageState tenantUsageStateMock; - @Mock private NotificationRuleProcessor notificationRuleProcessor; @@ -73,21 +74,27 @@ public class DefaultTbApiUsageStateServiceTest { @Mock private InternalTelemetryService tsWsService; + @Mock + private MailExecutorService mailExecutor; + @Mock private PartitionService partitionService; @Mock - private MailExecutorService mailExecutor; + private TbTenantProfileCache tbTenantProfileCache; @Spy @InjectMocks DefaultTbApiUsageStateService service; - private ApiUsageState dummyUsageState; private TenantApiUsageState tenantApiUsageState; private Tenant dummyTenant; private static final int MAX_ENABLE_VALUE = 5000; + private static final long VALUE_WARNING = 4500L; + private static final long VALUE_DISABLE = 5500L; + private static final int NEW_MAX_ENABLE_VALUE = 4000; + private static final double NEW_WARN_THRESHOLD_VALUE = 0.9; private static final double WARN_THRESHOLD_VALUE = 0.8; @BeforeEach @@ -117,11 +124,20 @@ public class DefaultTbApiUsageStateServiceTest { TenantProfileData profileData = new TenantProfileData(); profileData.setConfiguration(config); - TenantProfile tenantProfile = new TenantProfile(); - tenantProfile.setId(new TenantProfileId(UUID.randomUUID())); - tenantProfile.setProfileData(profileData); + TenantProfile dummyTenantProfile = new TenantProfile(); + dummyTenantProfile.setId(new TenantProfileId(UUID.randomUUID())); + dummyTenantProfile.setProfileData(profileData); - dummyUsageState = new ApiUsageState(); + ApiUsageState dummyUsageState = getApiUsageState(tenantId); + + tenantApiUsageState = new TenantApiUsageState(dummyTenantProfile, dummyUsageState); + + service.myUsageStates.put(tenantId, tenantApiUsageState); + } + + @NotNull + private static ApiUsageState getApiUsageState(TenantId tenantId) { + ApiUsageState dummyUsageState = new ApiUsageState(); dummyUsageState.setTransportState(ApiUsageStateValue.ENABLED); dummyUsageState.setDbStorageState(ApiUsageStateValue.ENABLED); dummyUsageState.setReExecState(ApiUsageStateValue.ENABLED); @@ -133,25 +149,63 @@ public class DefaultTbApiUsageStateServiceTest { dummyUsageState.setTenantId(tenantId); dummyUsageState.setEntityId(tenantId); dummyUsageState.setVersion(1L); - - tenantApiUsageState = new TenantApiUsageState(tenantProfile, dummyUsageState); - - service.myUsageStates.put(tenantId, tenantApiUsageState); - } - - @AfterEach - public void tearDown() { - + return dummyUsageState; } @Test public void givenTenantIdFromEntityStatesMap_whenGetApiUsageState() { - service.myUsageStates.put(dummyTenant.getTenantId(), tenantUsageStateMock); ApiUsageState tenantUsageState = service.getApiUsageState(dummyTenant.getTenantId()); - assertThat(tenantUsageState, is(tenantUsageStateMock.getApiUsageState())); + assertThat(tenantUsageState, is(tenantApiUsageState.getApiUsageState())); verify(service, never()).getOrFetchState(dummyTenant.getTenantId(), dummyTenant.getTenantId()); } + @Test + public void testGetApiUsageState_fromOtherUsageStates() { + TenantId tenantId = dummyTenant.getTenantId(); + service.myUsageStates.remove(tenantId); + ApiUsageState otherState = new ApiUsageState(); + otherState.setTenantId(tenantId); + service.otherUsageStates.put(tenantId, otherState); + ApiUsageState result = service.getApiUsageState(tenantId); + assertThat(result, is(otherState)); + } + + @Test + public void testGetApiUsageState_whenMyPartition() { + TenantId tenantId = dummyTenant.getTenantId(); + service.myUsageStates.remove(tenantId); + service.otherUsageStates.remove(tenantId); + + TopicPartitionInfo tpi = mock(TopicPartitionInfo.class); + when(tpi.isMyPartition()).thenReturn(true); + when(partitionService.resolve(any(), eq(dummyTenant.getId()), eq(dummyTenant.getId()))).thenReturn(tpi); + + TenantApiUsageState fetchedState = tenantApiUsageState; + doReturn(fetchedState).when(service).getOrFetchState(tenantId, tenantId); + ApiUsageState result = service.getApiUsageState(tenantId); + assertThat(result, is(fetchedState.getApiUsageState())); + } + + @Test + public void testGetApiUsageState_whenNotMyPartition() { + TenantId tenantId = dummyTenant.getTenantId(); + service.myUsageStates.remove(tenantId); + service.otherUsageStates.remove(tenantId); + + TopicPartitionInfo tpi = mock(TopicPartitionInfo.class); + when(tpi.isMyPartition()).thenReturn(false); + when(partitionService.resolve(any(), eq(tenantId), eq(tenantId))).thenReturn(tpi); + + ApiUsageState foundState = new ApiUsageState(); + foundState.setTenantId(tenantId); + when(apiUsageStateService.findTenantApiUsageState(tenantId)).thenReturn(foundState); + + ApiUsageState result = service.getApiUsageState(tenantId); + assertThat(result, is(foundState)); + + assertThat(service.otherUsageStates.get(tenantId), is(foundState)); + } + @Test public void testProcess_AllFeaturesTransitionToWarning() { TransportProtos.ToUsageStatsServiceMsg.Builder msgBuilder = TransportProtos.ToUsageStatsServiceMsg.newBuilder() @@ -165,7 +219,7 @@ public class DefaultTbApiUsageStateServiceTest { if (key.getApiFeature() != null) { msgBuilder.addValues(TransportProtos.UsageStatsKVProto.newBuilder() .setKey(key.name()) - .setValue(4500L) + .setValue(VALUE_WARNING) .build()); } } @@ -195,7 +249,6 @@ public class DefaultTbApiUsageStateServiceTest { assertEquals(ApiUsageStateValue.WARNING, tenantApiUsageState.getFeatureValue(ApiFeature.JS)); - verify(notificationRuleProcessor, atLeastOnce()).process(any()); verify(mailExecutor, atLeastOnce()).submit((Runnable) any()); } @@ -213,7 +266,7 @@ public class DefaultTbApiUsageStateServiceTest { if (key.getApiFeature() != null) { msgBuilder.addValues(TransportProtos.UsageStatsKVProto.newBuilder() .setKey(key.name()) - .setValue(5500L) + .setValue(VALUE_DISABLE) .build()); } } @@ -241,6 +294,73 @@ public class DefaultTbApiUsageStateServiceTest { verify(mailExecutor, atLeastOnce()).submit((Runnable) any()); } + @Test + public void testOnTenantProfileUpdate_updatesStateForMatchingTenant() { + TenantProfileId currentProfileId = tenantApiUsageState.getTenantProfileId(); + + TenantProfileData newProfileData = new TenantProfileData(); + DefaultTenantProfileConfiguration newConfig = DefaultTenantProfileConfiguration.builder() + .maxJSExecutions(NEW_MAX_ENABLE_VALUE) + .maxTransportMessages(MAX_ENABLE_VALUE) + .maxRuleChains(MAX_ENABLE_VALUE) + .maxTbelExecutions(MAX_ENABLE_VALUE) + .maxDPStorageDays(NEW_MAX_ENABLE_VALUE) + .maxREExecutions(MAX_ENABLE_VALUE) + .maxEmails(NEW_MAX_ENABLE_VALUE) + .maxSms(MAX_ENABLE_VALUE) + .maxCreatedAlarms(MAX_ENABLE_VALUE) + .warnThreshold(NEW_WARN_THRESHOLD_VALUE) + .build(); + + newProfileData.setConfiguration(newConfig); + + TenantProfile newProfile = new TenantProfile(); + newProfile.setId(currentProfileId); + newProfile.setProfileData(newProfileData); + + when(tbTenantProfileCache.get(currentProfileId)).thenReturn(newProfile); + + service.onTenantProfileUpdate(currentProfileId); + + assertEquals(currentProfileId, tenantApiUsageState.getTenantProfileId()); + assertEquals(newProfileData, tenantApiUsageState.getTenantProfileData()); + } + + @Test + public void testOnTenantUpdate_updatesStateWhenProfileChanged() { + TenantProfileId oldProfileId = tenantApiUsageState.getTenantProfileId(); + + TenantProfileId newProfileId = new TenantProfileId(UUID.randomUUID()); + TenantProfileData newProfileData = new TenantProfileData(); + DefaultTenantProfileConfiguration newConfig = DefaultTenantProfileConfiguration.builder() + .maxJSExecutions(NEW_MAX_ENABLE_VALUE) + .maxTransportMessages(MAX_ENABLE_VALUE) + .maxRuleChains(MAX_ENABLE_VALUE) + .maxTbelExecutions(NEW_MAX_ENABLE_VALUE) + .maxDPStorageDays(MAX_ENABLE_VALUE) + .maxREExecutions(NEW_MAX_ENABLE_VALUE) + .maxEmails(MAX_ENABLE_VALUE) + .maxSms(MAX_ENABLE_VALUE) + .maxCreatedAlarms(MAX_ENABLE_VALUE) + .warnThreshold(NEW_WARN_THRESHOLD_VALUE) + .build(); + + newProfileData.setConfiguration(newConfig); + + TenantProfile newProfile = new TenantProfile(); + newProfile.setId(newProfileId); + newProfile.setProfileData(newProfileData); + + when(tbTenantProfileCache.get(dummyTenant.getTenantId())).thenReturn(newProfile); + + assertNotEquals(oldProfileId, newProfileId); + + service.onTenantUpdate(dummyTenant.getTenantId()); + + assertEquals(newProfileId, tenantApiUsageState.getTenantProfileId()); + assertEquals(newProfileData, tenantApiUsageState.getTenantProfileData()); + } + private boolean containsFeature(ApiFeature feature) { for (ApiUsageRecordKey key : ApiUsageRecordKey.values()) { if (key.getApiFeature() != null && key.getApiFeature().equals(feature)) {