refactor & update logic

application/src/main/java/org/thingsboard/server/actors/ActorSystemContext.java

@@ -75,6 +75,8 @@ import org.thingsboard.server.dao.tenant.TenantProfileService;
 import org.thingsboard.server.dao.tenant.TenantService;
 import org.thingsboard.server.dao.timeseries.TimeseriesService;
 import org.thingsboard.server.dao.user.UserService;
+import org.thingsboard.server.dao.widget.WidgetTypeService;
+import org.thingsboard.server.dao.widget.WidgetsBundleService;
 import org.thingsboard.server.queue.discovery.PartitionService;
 import org.thingsboard.server.queue.discovery.TbServiceInfoProvider;
 import org.thingsboard.server.queue.usagestats.TbApiUsageClient;
@@ -377,6 +379,16 @@ public class ActorSystemContext {
     @Getter
     private QueueService queueService;
 
+    @Lazy
+    @Autowired(required = false)
+    @Getter
+    private WidgetsBundleService widgetsBundleService;
+
+    @Lazy
+    @Autowired(required = false)
+    @Getter
+    private WidgetTypeService widgetTypeService;
+
     @Value("${actors.session.max_concurrent_sessions_per_device:1}")
     @Getter
     private long maxConcurrentSessionsPerDevice;

application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java

@@ -25,6 +25,7 @@ import org.thingsboard.common.util.JacksonUtil;
 import org.thingsboard.common.util.ListeningExecutor;
 import org.thingsboard.rule.engine.api.MailService;
 import org.thingsboard.rule.engine.api.RuleEngineAlarmService;
+import org.thingsboard.rule.engine.api.RuleEngineApiUsageStateService;
 import org.thingsboard.rule.engine.api.RuleEngineAssetProfileCache;
 import org.thingsboard.rule.engine.api.RuleEngineDeviceProfileCache;
 import org.thingsboard.rule.engine.api.RuleEngineRpcService;
@@ -88,6 +89,8 @@ import org.thingsboard.server.dao.rule.RuleChainService;
 import org.thingsboard.server.dao.tenant.TenantService;
 import org.thingsboard.server.dao.timeseries.TimeseriesService;
 import org.thingsboard.server.dao.user.UserService;
+import org.thingsboard.server.dao.widget.WidgetTypeService;
+import org.thingsboard.server.dao.widget.WidgetsBundleService;
 import org.thingsboard.server.gen.transport.TransportProtos;
 import org.thingsboard.server.queue.TbQueueCallback;
 import org.thingsboard.server.queue.TbQueueMsgMetadata;
@@ -719,6 +722,26 @@ class DefaultTbContext implements TbContext {
         return mainCtx.getTenantProfileCache().get(getTenantId());
     }
 
+    @Override
+    public TenantProfile getTenantProfile(TenantId tenantId) {
+        return mainCtx.getTenantProfileCache().get(tenantId);
+    }
+
+    @Override
+    public WidgetsBundleService getWidgetBundleService() {
+        return mainCtx.getWidgetsBundleService();
+    }
+
+    @Override
+    public WidgetTypeService getWidgetTypeService() {
+        return mainCtx.getWidgetTypeService();
+    }
+
+    @Override
+    public RuleEngineApiUsageStateService getRuleEngineApiUsageStateService() {
+        return mainCtx.getApiUsageStateService();
+    }
+
     private TbMsgMetaData getActionMetaData(RuleNodeId ruleNodeId) {
         TbMsgMetaData metaData = new TbMsgMetaData();
         metaData.putValue("ruleNodeId", ruleNodeId.toString());

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

@@ -395,6 +395,11 @@ public class DefaultTbApiUsageStateService extends AbstractPartitionBasedService
         }
     }
 
+    @Override
+    public ApiUsageState findApiUsageStateById(TenantId tenantId, ApiUsageStateId id) {
+        return apiUsageStateService.findApiUsageStateById(tenantId, id);
+    }
+
     private interface StateChecker {
         boolean check(long threshold, long warnThreshold, long value);
     }

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

@@ -16,6 +16,7 @@
 package org.thingsboard.server.service.apiusage;
 
 import org.springframework.context.ApplicationListener;
+import org.thingsboard.rule.engine.api.RuleEngineApiUsageStateService;
 import org.thingsboard.server.common.data.ApiUsageState;
 import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.TenantId;
@@ -25,7 +26,7 @@ import org.thingsboard.server.gen.transport.TransportProtos.ToUsageStatsServiceM
 import org.thingsboard.server.queue.common.TbProtoQueueMsg;
 import org.thingsboard.server.queue.discovery.event.PartitionChangeEvent;
 
-public interface TbApiUsageStateService extends ApplicationListener<PartitionChangeEvent> {
+public interface TbApiUsageStateService extends RuleEngineApiUsageStateService, ApplicationListener<PartitionChangeEvent> {
 
     void process(TbProtoQueueMsg<ToUsageStatsServiceMsg> msg, TbCallback callback);
 

application/src/main/java/org/thingsboard/server/service/rpc/DefaultTbRuleEngineRpcService.java

@@ -15,10 +15,14 @@
  */
 package org.thingsboard.server.service.rpc;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 import org.thingsboard.common.util.ThingsBoardThreadFactory;
+import org.thingsboard.server.common.data.id.RpcId;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.data.rpc.Rpc;
 import org.thingsboard.server.common.data.rpc.RpcError;
 import org.thingsboard.rule.engine.api.RuleEngineDeviceRpcRequest;
 import org.thingsboard.rule.engine.api.RuleEngineDeviceRpcResponse;
@@ -27,6 +31,7 @@ import org.thingsboard.server.common.msg.queue.ServiceType;
 import org.thingsboard.server.common.msg.queue.TopicPartitionInfo;
 import org.thingsboard.server.common.msg.rpc.FromDeviceRpcResponse;
 import org.thingsboard.server.common.msg.rpc.ToDeviceRpcRequest;
+import org.thingsboard.server.dao.rpc.RpcService;
 import org.thingsboard.server.gen.transport.TransportProtos;
 import org.thingsboard.server.queue.discovery.PartitionService;
 import org.thingsboard.server.queue.discovery.TbServiceInfoProvider;
@@ -52,6 +57,7 @@ public class DefaultTbRuleEngineRpcService implements TbRuleEngineDeviceRpcServi
     private final PartitionService partitionService;
     private final TbClusterService clusterService;
     private final TbServiceInfoProvider serviceInfoProvider;
+    private final RpcService rpcService;
 
     private final ConcurrentMap<UUID, Consumer<FromDeviceRpcResponse>> toDeviceRpcRequests = new ConcurrentHashMap<>();
 
@@ -61,10 +67,11 @@ public class DefaultTbRuleEngineRpcService implements TbRuleEngineDeviceRpcServi
 
     public DefaultTbRuleEngineRpcService(PartitionService partitionService,
                                          TbClusterService clusterService,
-                                         TbServiceInfoProvider serviceInfoProvider) {
+                                         TbServiceInfoProvider serviceInfoProvider, RpcService rpcService) {
         this.partitionService = partitionService;
         this.clusterService = clusterService;
         this.serviceInfoProvider = serviceInfoProvider;
+        this.rpcService = rpcService;
     }
 
     @Autowired(required = false)
@@ -119,6 +126,11 @@ public class DefaultTbRuleEngineRpcService implements TbRuleEngineDeviceRpcServi
         });
     }
 
+    @Override
+    public ListenableFuture<Rpc> findRpcByIdAsync(TenantId tenantId, RpcId id) {
+        return rpcService.findRpcByIdAsync(tenantId, id);
+    }
+
     @Override
     public void processRpcResponseFromDevice(FromDeviceRpcResponse response) {
         log.trace("[{}] Received response to server-side RPC request from Core RPC Service", response.getId());

rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineApiUsageStateService.java

@@ -0,0 +1,26 @@
+/**
+ * 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.rule.engine.api;
+
+import org.thingsboard.server.common.data.ApiUsageState;
+import org.thingsboard.server.common.data.id.ApiUsageStateId;
+import org.thingsboard.server.common.data.id.TenantId;
+
+public interface RuleEngineApiUsageStateService {
+
+    ApiUsageState findApiUsageStateById(TenantId tenantId, ApiUsageStateId id);
+
+}

rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/RuleEngineRpcService.java

@@ -15,7 +15,11 @@
  */
 package org.thingsboard.rule.engine.api;
 
+import com.google.common.util.concurrent.ListenableFuture;
 import org.thingsboard.server.common.data.id.DeviceId;
+import org.thingsboard.server.common.data.id.RpcId;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.data.rpc.Rpc;
 
 import java.util.UUID;
 import java.util.function.Consumer;
@@ -29,4 +33,5 @@ public interface RuleEngineRpcService {
 
     void sendRpcRequestToDevice(RuleEngineDeviceRpcRequest request, Consumer<RuleEngineDeviceRpcResponse> consumer);
 
+    ListenableFuture<Rpc> findRpcByIdAsync(TenantId tenantId, RpcId id);
 }

rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java

@@ -61,6 +61,8 @@ import org.thingsboard.server.dao.rule.RuleChainService;
 import org.thingsboard.server.dao.tenant.TenantService;
 import org.thingsboard.server.dao.timeseries.TimeseriesService;
 import org.thingsboard.server.dao.user.UserService;
+import org.thingsboard.server.dao.widget.WidgetTypeService;
+import org.thingsboard.server.dao.widget.WidgetsBundleService;
 
 import java.util.List;
 import java.util.Set;
@@ -305,4 +307,12 @@ public interface TbContext {
     void removeListeners();
 
     TenantProfile getTenantProfile();
+
+    TenantProfile getTenantProfile(TenantId tenantId);
+
+    WidgetsBundleService getWidgetBundleService();
+
+    WidgetTypeService getWidgetTypeService();
+
+    RuleEngineApiUsageStateService getRuleEngineApiUsageStateService();
 }

rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/util/EntitiesTenantIdAsyncLoader.java

@@ -20,8 +20,10 @@ import com.google.common.util.concurrent.ListenableFuture;
 import org.thingsboard.common.util.ListeningExecutor;
 import org.thingsboard.rule.engine.api.TbContext;
 import org.thingsboard.rule.engine.api.TbNodeException;
+import org.thingsboard.server.common.data.EntityType;
 import org.thingsboard.server.common.data.HasTenantId;
 import org.thingsboard.server.common.data.id.AlarmId;
+import org.thingsboard.server.common.data.id.ApiUsageStateId;
 import org.thingsboard.server.common.data.id.AssetId;
 import org.thingsboard.server.common.data.id.AssetProfileId;
 import org.thingsboard.server.common.data.id.CustomerId;
@@ -32,50 +34,104 @@ import org.thingsboard.server.common.data.id.EdgeId;
 import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.id.EntityViewId;
 import org.thingsboard.server.common.data.id.OtaPackageId;
+import org.thingsboard.server.common.data.id.QueueId;
+import org.thingsboard.server.common.data.id.RpcId;
 import org.thingsboard.server.common.data.id.RuleChainId;
+import org.thingsboard.server.common.data.id.RuleNodeId;
+import org.thingsboard.server.common.data.id.TbResourceId;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.common.data.id.UserId;
+import org.thingsboard.server.common.data.id.WidgetTypeId;
+import org.thingsboard.server.common.data.id.WidgetsBundleId;
+import org.thingsboard.server.common.data.rule.RuleNode;
+
+import java.util.UUID;
 
 public class EntitiesTenantIdAsyncLoader {
 
     public static ListenableFuture<TenantId> findEntityIdAsync(TbContext ctx, EntityId original) {
-        ListeningExecutor executor = ctx.getDbCallbackExecutor();
+        ListenableFuture<? extends HasTenantId> hasTenantId;
-        switch (original.getEntityType()) {
+        UUID id = original.getId();
+        EntityType entityType = original.getEntityType();
+        TenantId tenantId = ctx.getTenantId();
+        switch (entityType) {
             case TENANT:
-                return Futures.immediateFuture((TenantId) original);
+                return Futures.immediateFuture(new TenantId(id));
             case CUSTOMER:
-                return getTenantAsync(ctx.getCustomerService().findCustomerByIdAsync(ctx.getTenantId(), (CustomerId) original), executor);
+                hasTenantId = ctx.getCustomerService().findCustomerByIdAsync(tenantId, new CustomerId(id));
+                break;
             case USER:
-                return getTenantAsync(ctx.getUserService().findUserByIdAsync(ctx.getTenantId(), (UserId) original), executor);
+                hasTenantId = ctx.getUserService().findUserByIdAsync(tenantId, new UserId(id));
+                break;
             case ASSET:
-                return getTenantAsync(ctx.getAssetService().findAssetByIdAsync(ctx.getTenantId(), (AssetId) original), executor);
+                hasTenantId = ctx.getAssetService().findAssetByIdAsync(tenantId, new AssetId(id));
+                break;
             case DEVICE:
-                return getTenantAsync(ctx.getDeviceService().findDeviceByIdAsync(ctx.getTenantId(), (DeviceId) original), executor);
+                hasTenantId = ctx.getDeviceService().findDeviceByIdAsync(tenantId, new DeviceId(id));
+                break;
             case ALARM:
-                return getTenantAsync(ctx.getAlarmService().findAlarmByIdAsync(ctx.getTenantId(), (AlarmId) original), executor);
+                hasTenantId = ctx.getAlarmService().findAlarmByIdAsync(tenantId, new AlarmId(id));
+                break;
             case RULE_CHAIN:
-                return getTenantAsync(ctx.getRuleChainService().findRuleChainByIdAsync(ctx.getTenantId(), (RuleChainId) original), executor);
+                hasTenantId = ctx.getRuleChainService().findRuleChainByIdAsync(tenantId, new RuleChainId(id));
+                break;
             case ENTITY_VIEW:
-                return getTenantAsync(ctx.getEntityViewService().findEntityViewByIdAsync(ctx.getTenantId(), (EntityViewId) original), executor);
+                hasTenantId = ctx.getEntityViewService().findEntityViewByIdAsync(tenantId, new EntityViewId(id));
+                break;
             case DASHBOARD:
-                return getTenantAsync(ctx.getDashboardService().findDashboardByIdAsync(ctx.getTenantId(), (DashboardId) original), executor);
+                hasTenantId = ctx.getDashboardService().findDashboardByIdAsync(tenantId, new DashboardId(id));
+                break;
             case EDGE:
-                return getTenantAsync(ctx.getEdgeService().findEdgeByIdAsync(ctx.getTenantId(), (EdgeId) original), executor);
+                hasTenantId = ctx.getEdgeService().findEdgeByIdAsync(tenantId, new EdgeId(id));
+                break;
             case OTA_PACKAGE:
-                return getTenantAsync(ctx.getOtaPackageService().findOtaPackageInfoByIdAsync(ctx.getTenantId(), (OtaPackageId) original), executor);
+                hasTenantId = ctx.getOtaPackageService().findOtaPackageInfoByIdAsync(tenantId, new OtaPackageId(id));
+                break;
             case ASSET_PROFILE:
-                return getTenantAsync(Futures.immediateFuture(ctx.getAssetProfileCache().get(ctx.getTenantId(), (AssetProfileId) original)), executor);
+                hasTenantId = Futures.immediateFuture(ctx.getAssetProfileCache().get(tenantId, new AssetProfileId(id)));
+                break;
             case DEVICE_PROFILE:
-                return getTenantAsync(Futures.immediateFuture(ctx.getDeviceProfileCache().get(ctx.getTenantId(), (DeviceProfileId) original)), executor);
+                hasTenantId = Futures.immediateFuture(ctx.getDeviceProfileCache().get(tenantId, new DeviceProfileId(id)));
+                break;
+            case WIDGET_TYPE:
+                hasTenantId = Futures.immediateFuture(ctx.getWidgetTypeService().findWidgetTypeById(tenantId, new WidgetTypeId(id)));
+                break;
+            case WIDGETS_BUNDLE:
+                hasTenantId = Futures.immediateFuture(ctx.getWidgetBundleService().findWidgetsBundleById(tenantId, new WidgetsBundleId(id)));
+                break;
+            case RPC:
+                hasTenantId = ctx.getRpcService().findRpcByIdAsync(ctx.getTenantId(), new RpcId(id));
+                break;
+            case QUEUE:
+                hasTenantId = Futures.immediateFuture(ctx.getQueueService().findQueueById(tenantId, new QueueId(id)));
+                break;
+            case API_USAGE_STATE:
+                hasTenantId = Futures.immediateFuture(ctx.getRuleEngineApiUsageStateService().findApiUsageStateById(tenantId, new ApiUsageStateId(id)));
+                break;
+            case TB_RESOURCE:
+                hasTenantId = ctx.getResourceService().findResourceInfoByIdAsync(tenantId, new TbResourceId(id));
+                break;
+            case RULE_NODE:
+                hasTenantId = null;
+                RuleNode ruleNode = ctx.getRuleChainService().findRuleNodeById(tenantId, new RuleNodeId(id));
+                if (ruleNode != null) {
+                    return Futures.immediateFuture(tenantId);
+                }
+                break;
+            case TENANT_PROFILE:
+                hasTenantId = null;
+                if (ctx.getTenantProfile() == ctx.getTenantProfile(new TenantId(id))) {
+                    return Futures.immediateFuture(tenantId);
+                }
+                break;
             default:
-                return Futures.immediateFailedFuture(new TbNodeException("Unexpected original EntityType " + original.getEntityType()));
+                hasTenantId = Futures.immediateFailedFuture(new TbNodeException("Unexpected original EntityType " + original.getEntityType()));
         }
+        return getTenantIdAsync(hasTenantId, ctx.getDbCallbackExecutor());
     }
 
-    private static <T extends HasTenantId> ListenableFuture<TenantId> getTenantAsync(ListenableFuture<T> future, ListeningExecutor executor) {
+    private static <T extends HasTenantId> ListenableFuture<TenantId> getTenantIdAsync(ListenableFuture<T> future, ListeningExecutor executor) {
-        return Futures.transformAsync(future, in -> {
+        return Futures.transformAsync(future, in -> in != null ? Futures.immediateFuture(in.getTenantId())
-            return in != null ? Futures.immediateFuture(in.getTenantId())
+                : Futures.immediateFuture(null), executor);
-                    : Futures.immediateFuture(null);
-        }, executor);
     }
 }

rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/util/EntitiesTenantIdAsyncLoaderTest.java

@@ -0,0 +1,347 @@
+/**
+ * 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.rule.engine.util;
+
+import com.google.common.util.concurrent.Futures;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+import org.thingsboard.common.util.AbstractListeningExecutor;
+import org.thingsboard.rule.engine.api.RuleEngineAlarmService;
+import org.thingsboard.rule.engine.api.RuleEngineApiUsageStateService;
+import org.thingsboard.rule.engine.api.RuleEngineAssetProfileCache;
+import org.thingsboard.rule.engine.api.RuleEngineDeviceProfileCache;
+import org.thingsboard.rule.engine.api.RuleEngineRpcService;
+import org.thingsboard.rule.engine.api.TbContext;
+import org.thingsboard.server.common.data.ApiUsageState;
+import org.thingsboard.server.common.data.Customer;
+import org.thingsboard.server.common.data.Dashboard;
+import org.thingsboard.server.common.data.Device;
+import org.thingsboard.server.common.data.DeviceProfile;
+import org.thingsboard.server.common.data.EntityType;
+import org.thingsboard.server.common.data.EntityView;
+import org.thingsboard.server.common.data.OtaPackage;
+import org.thingsboard.server.common.data.TbResource;
+import org.thingsboard.server.common.data.TenantProfile;
+import org.thingsboard.server.common.data.User;
+import org.thingsboard.server.common.data.alarm.Alarm;
+import org.thingsboard.server.common.data.asset.Asset;
+import org.thingsboard.server.common.data.asset.AssetProfile;
+import org.thingsboard.server.common.data.edge.Edge;
+import org.thingsboard.server.common.data.id.AssetProfileId;
+import org.thingsboard.server.common.data.id.DeviceProfileId;
+import org.thingsboard.server.common.data.id.EntityId;
+import org.thingsboard.server.common.data.id.EntityIdFactory;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.data.queue.Queue;
+import org.thingsboard.server.common.data.rpc.Rpc;
+import org.thingsboard.server.common.data.rule.RuleChain;
+import org.thingsboard.server.common.data.rule.RuleNode;
+import org.thingsboard.server.common.data.widget.WidgetType;
+import org.thingsboard.server.common.data.widget.WidgetsBundle;
+import org.thingsboard.server.dao.asset.AssetService;
+import org.thingsboard.server.dao.customer.CustomerService;
+import org.thingsboard.server.dao.dashboard.DashboardService;
+import org.thingsboard.server.dao.device.DeviceService;
+import org.thingsboard.server.dao.edge.EdgeService;
+import org.thingsboard.server.dao.entityview.EntityViewService;
+import org.thingsboard.server.dao.ota.OtaPackageService;
+import org.thingsboard.server.dao.queue.QueueService;
+import org.thingsboard.server.dao.resource.ResourceService;
+import org.thingsboard.server.dao.rule.RuleChainService;
+import org.thingsboard.server.dao.user.UserService;
+import org.thingsboard.server.dao.widget.WidgetTypeService;
+import org.thingsboard.server.dao.widget.WidgetsBundleService;
+
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class EntitiesTenantIdAsyncLoaderTest {
+
+    @Mock
+    private TbContext ctx;
+    @Mock
+    private CustomerService customerService;
+    @Mock
+    private UserService userService;
+    @Mock
+    private AssetService assetService;
+    @Mock
+    private DeviceService deviceService;
+    @Mock
+    private RuleEngineAlarmService alarmService;
+    @Mock
+    private RuleChainService ruleChainService;
+    @Mock
+    private EntityViewService entityViewService;
+    @Mock
+    private DashboardService dashboardService;
+    @Mock
+    private EdgeService edgeService;
+    @Mock
+    private OtaPackageService otaPackageService;
+    @Mock
+    private RuleEngineAssetProfileCache assetProfileCache;
+    @Mock
+    private RuleEngineDeviceProfileCache deviceProfileCache;
+    @Mock
+    private WidgetTypeService widgetTypeService;
+    @Mock
+    private WidgetsBundleService widgetsBundleService;
+    @Mock
+    private QueueService queueService;
+    @Mock
+    private ResourceService resourceService;
+    @Mock
+    private RuleEngineRpcService rpcService;
+    @Mock
+    private RuleEngineApiUsageStateService ruleEngineApiUsageStateService;
+
+    private TenantId tenantId;
+    private AbstractListeningExecutor dbExecutor;
+
+    @Before
+    public void before() {
+        dbExecutor = new AbstractListeningExecutor() {
+            @Override
+            protected int getThreadPollSize() {
+                return 3;
+            }
+        };
+        dbExecutor.init();
+        this.tenantId = new TenantId(UUID.randomUUID());
+
+        when(ctx.getTenantId()).thenReturn(tenantId);
+        when(ctx.getDbCallbackExecutor()).thenReturn(dbExecutor);
+
+        for (EntityType entityType : EntityType.values()) {
+            initMocks(entityType, tenantId);
+        }
+    }
+
+    @After
+    public void after() {
+        dbExecutor.destroy();
+    }
+
+    private void initMocks(EntityType entityType, TenantId tenantId) {
+        switch (entityType) {
+            case TENANT:
+                break;
+            case CUSTOMER:
+                Customer customer = new Customer();
+                customer.setTenantId(tenantId);
+
+                when(ctx.getCustomerService()).thenReturn(customerService);
+                doReturn(Futures.immediateFuture(customer)).when(customerService).findCustomerByIdAsync(eq(tenantId), any());
+
+                break;
+            case USER:
+                User user = new User();
+                user.setTenantId(tenantId);
+
+                when(ctx.getUserService()).thenReturn(userService);
+                doReturn(Futures.immediateFuture(user)).when(userService).findUserByIdAsync(eq(tenantId), any());
+
+                break;
+            case ASSET:
+                Asset asset = new Asset();
+                asset.setTenantId(tenantId);
+
+                when(ctx.getAssetService()).thenReturn(assetService);
+                doReturn(Futures.immediateFuture(asset)).when(assetService).findAssetByIdAsync(eq(tenantId), any());
+
+                break;
+            case DEVICE:
+                Device device = new Device();
+                device.setTenantId(tenantId);
+
+                when(ctx.getDeviceService()).thenReturn(deviceService);
+                doReturn(Futures.immediateFuture(device)).when(deviceService).findDeviceByIdAsync(eq(tenantId), any());
+
+                break;
+            case ALARM:
+                Alarm alarm = new Alarm();
+                alarm.setTenantId(tenantId);
+
+                when(ctx.getAlarmService()).thenReturn(alarmService);
+                doReturn(Futures.immediateFuture(alarm)).when(alarmService).findAlarmByIdAsync(eq(tenantId), any());
+
+                break;
+            case RULE_CHAIN:
+                RuleChain ruleChain = new RuleChain();
+                ruleChain.setTenantId(tenantId);
+
+                when(ctx.getRuleChainService()).thenReturn(ruleChainService);
+                doReturn(Futures.immediateFuture(ruleChain)).when(ruleChainService).findRuleChainByIdAsync(eq(tenantId), any());
+
+                break;
+            case ENTITY_VIEW:
+                EntityView entityView = new EntityView();
+                entityView.setTenantId(tenantId);
+
+                when(ctx.getEntityViewService()).thenReturn(entityViewService);
+                doReturn(Futures.immediateFuture(entityView)).when(entityViewService).findEntityViewByIdAsync(eq(tenantId), any());
+
+                break;
+            case DASHBOARD:
+                Dashboard dashboard = new Dashboard();
+                dashboard.setTenantId(tenantId);
+
+                when(ctx.getDashboardService()).thenReturn(dashboardService);
+                doReturn(Futures.immediateFuture(dashboard)).when(dashboardService).findDashboardByIdAsync(eq(tenantId), any());
+
+                break;
+            case EDGE:
+                Edge edge = new Edge();
+                edge.setTenantId(tenantId);
+
+                when(ctx.getEdgeService()).thenReturn(edgeService);
+                doReturn(Futures.immediateFuture(edge)).when(edgeService).findEdgeByIdAsync(eq(tenantId), any());
+
+                break;
+            case OTA_PACKAGE:
+                OtaPackage otaPackage = new OtaPackage();
+                otaPackage.setTenantId(tenantId);
+
+                when(ctx.getOtaPackageService()).thenReturn(otaPackageService);
+                doReturn(Futures.immediateFuture(otaPackage)).when(otaPackageService).findOtaPackageInfoByIdAsync(eq(tenantId), any());
+
+                break;
+            case ASSET_PROFILE:
+                AssetProfile assetProfile = new AssetProfile();
+                assetProfile.setTenantId(tenantId);
+
+                when(ctx.getAssetProfileCache()).thenReturn(assetProfileCache);
+                doReturn(assetProfile).when(assetProfileCache).get(eq(tenantId), any(AssetProfileId.class));
+
+                break;
+            case DEVICE_PROFILE:
+                DeviceProfile deviceProfile = new DeviceProfile();
+                deviceProfile.setTenantId(tenantId);
+
+                when(ctx.getDeviceProfileCache()).thenReturn(deviceProfileCache);
+                doReturn(deviceProfile).when(deviceProfileCache).get(eq(tenantId), any(DeviceProfileId.class));
+
+                break;
+            case WIDGET_TYPE:
+                WidgetType widgetType = new WidgetType();
+                widgetType.setTenantId(tenantId);
+
+                when(ctx.getWidgetTypeService()).thenReturn(widgetTypeService);
+                doReturn(widgetType).when(widgetTypeService).findWidgetTypeById(eq(tenantId), any());
+
+                break;
+            case WIDGETS_BUNDLE:
+                WidgetsBundle widgetsBundle = new WidgetsBundle();
+                widgetsBundle.setTenantId(tenantId);
+
+                when(ctx.getWidgetBundleService()).thenReturn(widgetsBundleService);
+                doReturn(widgetsBundle).when(widgetsBundleService).findWidgetsBundleById(eq(tenantId), any());
+
+                break;
+            case RPC:
+                Rpc rps = new Rpc();
+                rps.setTenantId(tenantId);
+
+                when(ctx.getRpcService()).thenReturn(rpcService);
+                doReturn(Futures.immediateFuture(rps)).when(rpcService).findRpcByIdAsync(eq(tenantId), any());
+
+                break;
+            case QUEUE:
+                Queue queue = new Queue();
+                queue.setTenantId(tenantId);
+
+                when(ctx.getQueueService()).thenReturn(queueService);
+                doReturn(queue).when(queueService).findQueueById(eq(tenantId), any());
+
+                break;
+            case API_USAGE_STATE:
+                ApiUsageState apiUsageState = new ApiUsageState();
+                apiUsageState.setTenantId(tenantId);
+
+                when(ctx.getRuleEngineApiUsageStateService()).thenReturn(ruleEngineApiUsageStateService);
+                doReturn(apiUsageState).when(ruleEngineApiUsageStateService).findApiUsageStateById(eq(tenantId), any());
+
+                break;
+            case TB_RESOURCE:
+                TbResource tbResource = new TbResource();
+                tbResource.setTenantId(tenantId);
+
+                when(ctx.getResourceService()).thenReturn(resourceService);
+                doReturn(Futures.immediateFuture(tbResource)).when(resourceService).findResourceInfoByIdAsync(eq(tenantId), any());
+
+                break;
+            case RULE_NODE:
+                RuleNode ruleNode = new RuleNode();
+
+                when(ctx.getRuleChainService()).thenReturn(ruleChainService);
+                doReturn(ruleNode).when(ruleChainService).findRuleNodeById(eq(tenantId), any());
+
+                break;
+            case TENANT_PROFILE:
+                TenantProfile tenantProfile = new TenantProfile();
+
+                when(ctx.getTenantProfile()).thenReturn(tenantProfile);
+                when(ctx.getTenantProfile(any(TenantId.class))).thenReturn(tenantProfile);
+
+                break;
+            default:
+                throw new RuntimeException("Unexpected original EntityType " + entityType);
+        }
+
+    }
+
+    private EntityId getEntityId(EntityType entityType) {
+        return EntityIdFactory.getByTypeAndUuid(entityType, UUID.randomUUID());
+    }
+
+    private void checkTenant(TenantId checkTenantId, boolean equals) throws ExecutionException, InterruptedException {
+        for (EntityType entityType : EntityType.values()) {
+            EntityId entityId = EntityType.TENANT.equals(entityType) ? tenantId : getEntityId(entityType);
+            TenantId targetTenantId = EntitiesTenantIdAsyncLoader.findEntityIdAsync(ctx, entityId).get();
+
+            Assert.assertNotNull(targetTenantId);
+            String msg = "Check entity type <" + entityType.name() + ">:";
+            if (equals) {
+                Assert.assertEquals(msg, targetTenantId, checkTenantId);
+            } else {
+                Assert.assertNotEquals(msg, targetTenantId, checkTenantId);
+            }
+        }
+    }
+
+    @Test
+    public void test_findEntityIdAsync_current_tenant() throws ExecutionException, InterruptedException {
+        checkTenant(tenantId, true);
+    }
+
+    @Test
+    public void test_findEntityIdAsync_other_tenant() throws ExecutionException, InterruptedException {
+        checkTenant(new TenantId(UUID.randomUUID()), false);
+    }
+
+}