added tests for DefaultTbApiUsageStateService

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

@@ -120,13 +120,13 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
     private InternalTelemetryService tsWsService;
 
     // Entities that should be processed on this server
-    private final Map<EntityId, BaseApiUsageState> myUsageStates = new ConcurrentHashMap<>();
+    final Map<EntityId, BaseApiUsageState> myUsageStates = new ConcurrentHashMap<>();
     // Entities that should be processed on other servers
-    private final Map<EntityId, ApiUsageState> otherUsageStates = new ConcurrentHashMap<>();
+    final Map<EntityId, ApiUsageState> otherUsageStates = new ConcurrentHashMap<>();
 
-    private final ConcurrentMap<TopicPartitionInfo, Set<EntityId>> partitionedEntities = new ConcurrentHashMap<>();
+    final ConcurrentMap<TopicPartitionInfo, Set<EntityId>> partitionedEntities = new ConcurrentHashMap<>();
 
-    private final Set<EntityId> deletedEntities = Collections.newSetFromMap(new ConcurrentHashMap<>());
+    final Set<EntityId> deletedEntities = Collections.newSetFromMap(new ConcurrentHashMap<>());
 
     @Value("${usage.stats.report.enabled:true}")
     private boolean enabled;
@@ -489,7 +489,7 @@ public class DefaultTbApiUsageStateService extends TbApplicationEventListener<Pa
         tsWsService.saveAndNotifyInternal(state.getTenantId(), state.getApiUsageState().getId(), counts, VOID_CALLBACK);
     }
 
-    private BaseApiUsageState getOrFetchState(TenantId tenantId, EntityId entityId) {
+    BaseApiUsageState getOrFetchState(TenantId tenantId, EntityId entityId) {
         if (entityId == null || entityId.isNullUid()) {
             entityId = tenantId;
         }

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

@@ -0,0 +1,101 @@
+/**
+ * Copyright © 2016-2022 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.apiusage;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.thingsboard.rule.engine.api.MailService;
+import org.thingsboard.server.cluster.TbClusterService;
+import org.thingsboard.server.common.data.ApiUsageState;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.msg.queue.ServiceType;
+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.timeseries.TimeseriesService;
+import org.thingsboard.server.dao.usagerecord.ApiUsageStateService;
+import org.thingsboard.server.queue.discovery.PartitionService;
+import org.thingsboard.server.queue.scheduler.SchedulerComponent;
+import org.thingsboard.server.service.executors.DbCallbackExecutorService;
+
+import java.util.UUID;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.BDDMockito.willReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+
+@RunWith(MockitoJUnitRunner.class)
+public class DefaultTbApiUsageStateServiceTest {
+
+    @Mock
+    TenantService tenantService;
+    @Mock
+    TimeseriesService tsService;
+    @Mock
+    TbClusterService clusterService;
+    @Mock
+    PartitionService partitionService;
+    @Mock
+    TenantApiUsageState tenantUsageStateMock;
+    @Mock
+    ApiUsageStateService apiUsageStateService;
+    @Mock
+    SchedulerComponent scheduler;
+    @Mock
+    TbTenantProfileCache tenantProfileCache;
+    @Mock
+    MailService mailService;
+    @Mock
+    DbCallbackExecutorService dbExecutor;
+
+    TenantId tenantId = TenantId.fromUUID(UUID.fromString("00797a3b-7aeb-4b5b-b57a-c2a810d0f112"));
+
+    DefaultTbApiUsageStateService service;
+
+    @Before
+    public void setUp() {
+        service = spy(new DefaultTbApiUsageStateService(clusterService, partitionService, tenantService, tsService, apiUsageStateService, scheduler, tenantProfileCache, mailService, dbExecutor));
+    }
+
+    @Test
+    public void givenTenantIdFromEntityStatesMap_whenGetApiUsageState() {
+        service.myUsageStates.put(tenantId, tenantUsageStateMock);
+        ApiUsageState tenantUsageState = service.getApiUsageState(tenantId);
+        assertThat(tenantUsageState, is(tenantUsageStateMock.getApiUsageState()));
+        Mockito.verify(service, never()).getOrFetchState(tenantId, tenantId);
+    }
+
+    @Test
+    public void givenTenantIdWithoutTenantStateInMap_whenGetState_thenGetOrFetchState() {
+        TopicPartitionInfo tpi = Mockito.mock(TopicPartitionInfo.class);
+        Mockito.when(tpi.isMyPartition()).thenReturn(true);
+        willReturn(tpi).given(partitionService).resolve(ServiceType.TB_CORE, tenantId, tenantId);
+        service.myUsageStates.clear();
+        willReturn(tenantUsageStateMock).given(service).getOrFetchState(tenantId, tenantId);
+        ApiUsageState tenantUsageState = service.getApiUsageState(tenantId);
+        assertThat(tenantUsageState, is(tenantUsageStateMock.getApiUsageState()));
+        assertThat(true, is(service.partitionedEntities.get(tpi).contains(tenantId)));
+        Mockito.verify(service, times(1)).getOrFetchState(tenantId, tenantId);
+    }
+
+}