Tenant Id Loader fixes

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

@@ -5,7 +5,7 @@
  * 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
+ * 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,
@@ -36,7 +36,7 @@ import org.thingsboard.rule.engine.api.SmsService;
 import org.thingsboard.rule.engine.api.TbContext;
 import org.thingsboard.rule.engine.api.TbRelationTypes;
 import org.thingsboard.rule.engine.api.sms.SmsSenderFactory;
-import org.thingsboard.rule.engine.util.EntitiesTenantIdAsyncLoader;
+import org.thingsboard.rule.engine.util.TenantIdLoader;
 import org.thingsboard.server.actors.ActorSystemContext;
 import org.thingsboard.server.actors.TbActorRef;
 import org.thingsboard.server.cluster.TbClusterService;
@@ -776,17 +776,10 @@ class DefaultTbContext implements TbContext {
         return metaData;
     }
 
-    public void isTenantEntity(EntityId entityId) {
-        try {
-            TenantId entityTenantId = EntitiesTenantIdAsyncLoader.findEntityIdAsync(this, entityId).get();
-            if (entityTenantId == null) {
-                throw new RuntimeException("Entity with id '" + entityId + "'not found!");
-            }
-            if (!entityTenantId.equals(this.getTenantId())) {
-                throw new RuntimeException("Entity with id: '" + entityId + "' specified in the configuration doesn't belong to the current tenant.");
-            }
-        } catch (Exception e) {
-            throw new RuntimeException(e);
+    @Override
+    public void checkTenantEntity(EntityId entityId) {
+        if (!this.getTenantId().equals(TenantIdLoader.findTenantId(this, entityId))) {
+            throw new RuntimeException("Entity with id: '" + entityId + "' specified in the configuration doesn't belong to the current tenant.");
         }
     }
 

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

@@ -127,8 +127,8 @@ public class DefaultTbRuleEngineRpcService implements TbRuleEngineDeviceRpcServi
     }
 
     @Override
-    public ListenableFuture<Rpc> findRpcByIdAsync(TenantId tenantId, RpcId id) {
-        return rpcService.findRpcByIdAsync(tenantId, id);
+    public Rpc findRpcById(TenantId tenantId, RpcId id) {
+        return rpcService.findById(tenantId, id);
     }
 
     @Override

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

@@ -51,6 +51,8 @@ public interface RuleEngineAlarmService {
 
     ListenableFuture<Alarm> findAlarmByIdAsync(TenantId tenantId, AlarmId alarmId);
 
+    Alarm findAlarmById(TenantId tenantId, AlarmId alarmId);
+
     ListenableFuture<Alarm> findLatestByOriginatorAndType(TenantId tenantId, EntityId originator, String type);
 
     ListenableFuture<AlarmInfo> findAlarmInfoByIdAsync(TenantId tenantId, AlarmId alarmId);

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

@@ -33,5 +33,5 @@ public interface RuleEngineRpcService {
 
     void sendRpcRequestToDevice(RuleEngineDeviceRpcRequest request, Consumer<RuleEngineDeviceRpcResponse> consumer);
 
-    ListenableFuture<Rpc> findRpcByIdAsync(TenantId tenantId, RpcId id);
+    Rpc findRpcById(TenantId tenantId, RpcId id);
 }

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

@@ -199,6 +199,8 @@ public interface TbContext {
      *
      */
 
+    void checkTenantEntity(EntityId entityId);
+
     boolean isLocalEntity(EntityId entityId);
 
     RuleNodeId getSelfId();

rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/debug/TbMsgGeneratorNode.java

@@ -78,7 +78,7 @@ public class TbMsgGeneratorNode implements TbNode {
         this.currentMsgCount = 0;
         if (!StringUtils.isEmpty(config.getOriginatorId())) {
             originatorId = EntityIdFactory.getByTypeAndUuid(config.getOriginatorType(), config.getOriginatorId());
-            ctx.isTenantEntity(originatorId);
+            ctx.checkTenantEntity(originatorId);
         } else {
             originatorId = ctx.getSelfId();
         }

rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/filter/TbCheckRelationNode.java

@@ -62,7 +62,7 @@ public class TbCheckRelationNode implements TbNode {
         this.config = TbNodeUtils.convert(configuration, TbCheckRelationNodeConfiguration.class);
         if (config.isCheckForSingleEntity()) {
             this.singleEntityId = EntityIdFactory.getByTypeAndId(config.getEntityType(), config.getEntityId());
-            ctx.isTenantEntity(singleEntityId);
+            ctx.checkTenantEntity(singleEntityId);
         }
     }
 

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

@@ -1,139 +0,0 @@
-/**
- * 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 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;
-import org.thingsboard.server.common.data.id.DashboardId;
-import org.thingsboard.server.common.data.id.DeviceId;
-import org.thingsboard.server.common.data.id.DeviceProfileId;
-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) {
-        ListenableFuture<? extends HasTenantId> hasTenantId;
-        UUID id = original.getId();
-        EntityType entityType = original.getEntityType();
-        TenantId tenantId = ctx.getTenantId();
-        switch (entityType) {
-            case TENANT:
-                return Futures.immediateFuture(new TenantId(id));
-            case CUSTOMER:
-                hasTenantId = ctx.getCustomerService().findCustomerByIdAsync(tenantId, new CustomerId(id));
-                break;
-            case USER:
-                hasTenantId = ctx.getUserService().findUserByIdAsync(tenantId, new UserId(id));
-                break;
-            case ASSET:
-                hasTenantId = ctx.getAssetService().findAssetByIdAsync(tenantId, new AssetId(id));
-                break;
-            case DEVICE:
-                hasTenantId = ctx.getDeviceService().findDeviceByIdAsync(tenantId, new DeviceId(id));
-                break;
-            case ALARM:
-                hasTenantId = ctx.getAlarmService().findAlarmByIdAsync(tenantId, new AlarmId(id));
-                break;
-            case RULE_CHAIN:
-                hasTenantId = ctx.getRuleChainService().findRuleChainByIdAsync(tenantId, new RuleChainId(id));
-                break;
-            case ENTITY_VIEW:
-                hasTenantId = ctx.getEntityViewService().findEntityViewByIdAsync(tenantId, new EntityViewId(id));
-                break;
-            case DASHBOARD:
-                hasTenantId = ctx.getDashboardService().findDashboardByIdAsync(tenantId, new DashboardId(id));
-                break;
-            case EDGE:
-                hasTenantId = ctx.getEdgeService().findEdgeByIdAsync(tenantId, new EdgeId(id));
-                break;
-            case OTA_PACKAGE:
-                hasTenantId = ctx.getOtaPackageService().findOtaPackageInfoByIdAsync(tenantId, new OtaPackageId(id));
-                break;
-            case ASSET_PROFILE:
-                hasTenantId = Futures.immediateFuture(ctx.getAssetProfileCache().get(tenantId, new AssetProfileId(id)));
-                break;
-            case DEVICE_PROFILE:
-                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(tenantId, 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:
-                RuleNode ruleNode = ctx.getRuleChainService().findRuleNodeById(tenantId, new RuleNodeId(id));
-                if (ruleNode != null) {
-                    hasTenantId = ctx.getRuleChainService().findRuleChainByIdAsync(tenantId, ruleNode.getRuleChainId());
-                } else {
-                    hasTenantId = Futures.immediateFuture(null);
-                }
-                break;
-            case TENANT_PROFILE:
-                if (ctx.getTenantProfile().getId().equals(id)) {
-                    return Futures.immediateFuture(tenantId);
-                } else {
-                    hasTenantId = Futures.immediateFuture(null);
-                }
-                break;
-            default:
-                hasTenantId = Futures.immediateFailedFuture(new TbNodeException("Unexpected original EntityType " + original.getEntityType()));
-        }
-        return getTenantIdAsync(hasTenantId, ctx.getDbCallbackExecutor());
-    }
-
-    private static <T extends HasTenantId> ListenableFuture<TenantId> getTenantIdAsync(ListenableFuture<T> future, ListeningExecutor executor) {
-        return Futures.transformAsync(future, in -> in != null ? Futures.immediateFuture(in.getTenantId())
-                : Futures.immediateFuture(null), executor);
-    }
-}

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

@@ -0,0 +1,132 @@
+/**
+ * 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 org.thingsboard.rule.engine.api.TbContext;
+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;
+import org.thingsboard.server.common.data.id.DashboardId;
+import org.thingsboard.server.common.data.id.DeviceId;
+import org.thingsboard.server.common.data.id.DeviceProfileId;
+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 TenantIdLoader {
+
+    public static TenantId findTenantId(TbContext ctx, EntityId entityId) {
+        UUID id = entityId.getId();
+        EntityType entityType = entityId.getEntityType();
+        TenantId ctxTenantId = ctx.getTenantId();
+
+        HasTenantId tenantEntity;
+        switch (entityType) {
+            case TENANT:
+                return new TenantId(id);
+            case CUSTOMER:
+                tenantEntity = ctx.getCustomerService().findCustomerById(ctxTenantId, new CustomerId(id));
+                break;
+            case USER:
+                tenantEntity = ctx.getUserService().findUserById(ctxTenantId, new UserId(id));
+                break;
+            case ASSET:
+                tenantEntity = ctx.getAssetService().findAssetById(ctxTenantId, new AssetId(id));
+                break;
+            case DEVICE:
+                tenantEntity = ctx.getDeviceService().findDeviceById(ctxTenantId, new DeviceId(id));
+                break;
+            case ALARM:
+                tenantEntity = ctx.getAlarmService().findAlarmById(ctxTenantId, new AlarmId(id));
+                break;
+            case RULE_CHAIN:
+                tenantEntity = ctx.getRuleChainService().findRuleChainById(ctxTenantId, new RuleChainId(id));
+                break;
+            case ENTITY_VIEW:
+                tenantEntity = ctx.getEntityViewService().findEntityViewById(ctxTenantId, new EntityViewId(id));
+                break;
+            case DASHBOARD:
+                tenantEntity = ctx.getDashboardService().findDashboardById(ctxTenantId, new DashboardId(id));
+                break;
+            case EDGE:
+                tenantEntity = ctx.getEdgeService().findEdgeById(ctxTenantId, new EdgeId(id));
+                break;
+            case OTA_PACKAGE:
+                tenantEntity = ctx.getOtaPackageService().findOtaPackageInfoById(ctxTenantId, new OtaPackageId(id));
+                break;
+            case ASSET_PROFILE:
+                tenantEntity = ctx.getAssetProfileCache().get(ctxTenantId, new AssetProfileId(id));
+                break;
+            case DEVICE_PROFILE:
+                tenantEntity = ctx.getDeviceProfileCache().get(ctxTenantId, new DeviceProfileId(id));
+                break;
+            case WIDGET_TYPE:
+                tenantEntity = ctx.getWidgetTypeService().findWidgetTypeById(ctxTenantId, new WidgetTypeId(id));
+                break;
+            case WIDGETS_BUNDLE:
+                tenantEntity = ctx.getWidgetBundleService().findWidgetsBundleById(ctxTenantId, new WidgetsBundleId(id));
+                break;
+            case RPC:
+                tenantEntity = ctx.getRpcService().findRpcById(ctxTenantId, new RpcId(id));
+                break;
+            case QUEUE:
+                tenantEntity = ctx.getQueueService().findQueueById(ctxTenantId, new QueueId(id));
+                break;
+            case API_USAGE_STATE:
+                tenantEntity = ctx.getRuleEngineApiUsageStateService().findApiUsageStateById(ctxTenantId, new ApiUsageStateId(id));
+                break;
+            case TB_RESOURCE:
+                tenantEntity = ctx.getResourceService().findResourceInfoById(ctxTenantId, new TbResourceId(id));
+                break;
+            case RULE_NODE:
+                RuleNode ruleNode = ctx.getRuleChainService().findRuleNodeById(ctxTenantId, new RuleNodeId(id));
+                if (ruleNode != null) {
+                    tenantEntity = ctx.getRuleChainService().findRuleChainById(ctxTenantId, ruleNode.getRuleChainId());
+                } else {
+                    tenantEntity = null;
+                }
+                break;
+            case TENANT_PROFILE:
+                if (ctx.getTenantProfile().getId().equals(entityId)) {
+                    return ctxTenantId;
+                } else {
+                    tenantEntity = null;
+                }
+                break;
+            default:
+                throw new RuntimeException("Unexpected entity type: " + entityId.getEntityType());
+        }
+        return tenantEntity != null ? tenantEntity.getTenantId() : null;
+    }
+
+}

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

@@ -15,7 +15,6 @@
  */
 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;
@@ -50,6 +49,7 @@ 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.id.TenantProfileId;
 import org.thingsboard.server.common.data.queue.Queue;
 import org.thingsboard.server.common.data.rpc.Rpc;
 import org.thingsboard.server.common.data.rule.RuleChain;
@@ -71,7 +71,6 @@ 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;
@@ -79,7 +78,7 @@ import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.when;
 
 @RunWith(MockitoJUnitRunner.class)
-public class EntitiesTenantIdAsyncLoaderTest {
+public class TenantIdLoaderTest {
 
     @Mock
     private TbContext ctx;
@@ -121,6 +120,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
     private RuleEngineApiUsageStateService ruleEngineApiUsageStateService;
 
     private TenantId tenantId;
+    private TenantProfileId tenantProfileId;
     private AbstractListeningExecutor dbExecutor;
 
     @Before
@@ -133,9 +133,9 @@ public class EntitiesTenantIdAsyncLoaderTest {
         };
         dbExecutor.init();
         this.tenantId = new TenantId(UUID.randomUUID());
+        this.tenantProfileId = new TenantProfileId(UUID.randomUUID());
 
         when(ctx.getTenantId()).thenReturn(tenantId);
-        when(ctx.getDbCallbackExecutor()).thenReturn(dbExecutor);
 
         for (EntityType entityType : EntityType.values()) {
             initMocks(entityType, tenantId);
@@ -156,7 +156,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 customer.setTenantId(tenantId);
 
                 when(ctx.getCustomerService()).thenReturn(customerService);
-                doReturn(Futures.immediateFuture(customer)).when(customerService).findCustomerByIdAsync(eq(tenantId), any());
+                doReturn(customer).when(customerService).findCustomerById(eq(tenantId), any());
 
                 break;
             case USER:
@@ -164,7 +164,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 user.setTenantId(tenantId);
 
                 when(ctx.getUserService()).thenReturn(userService);
-                doReturn(Futures.immediateFuture(user)).when(userService).findUserByIdAsync(eq(tenantId), any());
+                doReturn(user).when(userService).findUserById(eq(tenantId), any());
 
                 break;
             case ASSET:
@@ -172,7 +172,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 asset.setTenantId(tenantId);
 
                 when(ctx.getAssetService()).thenReturn(assetService);
-                doReturn(Futures.immediateFuture(asset)).when(assetService).findAssetByIdAsync(eq(tenantId), any());
+                doReturn(asset).when(assetService).findAssetById(eq(tenantId), any());
 
                 break;
             case DEVICE:
@@ -180,7 +180,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 device.setTenantId(tenantId);
 
                 when(ctx.getDeviceService()).thenReturn(deviceService);
-                doReturn(Futures.immediateFuture(device)).when(deviceService).findDeviceByIdAsync(eq(tenantId), any());
+                doReturn(device).when(deviceService).findDeviceById(eq(tenantId), any());
 
                 break;
             case ALARM:
@@ -188,7 +188,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 alarm.setTenantId(tenantId);
 
                 when(ctx.getAlarmService()).thenReturn(alarmService);
-                doReturn(Futures.immediateFuture(alarm)).when(alarmService).findAlarmByIdAsync(eq(tenantId), any());
+                doReturn(alarm).when(alarmService).findAlarmById(eq(tenantId), any());
 
                 break;
             case RULE_CHAIN:
@@ -196,7 +196,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 ruleChain.setTenantId(tenantId);
 
                 when(ctx.getRuleChainService()).thenReturn(ruleChainService);
-                doReturn(Futures.immediateFuture(ruleChain)).when(ruleChainService).findRuleChainByIdAsync(eq(tenantId), any());
+                doReturn(ruleChain).when(ruleChainService).findRuleChainById(eq(tenantId), any());
 
                 break;
             case ENTITY_VIEW:
@@ -204,7 +204,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 entityView.setTenantId(tenantId);
 
                 when(ctx.getEntityViewService()).thenReturn(entityViewService);
-                doReturn(Futures.immediateFuture(entityView)).when(entityViewService).findEntityViewByIdAsync(eq(tenantId), any());
+                doReturn(entityView).when(entityViewService).findEntityViewById(eq(tenantId), any());
 
                 break;
             case DASHBOARD:
@@ -212,7 +212,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 dashboard.setTenantId(tenantId);
 
                 when(ctx.getDashboardService()).thenReturn(dashboardService);
-                doReturn(Futures.immediateFuture(dashboard)).when(dashboardService).findDashboardByIdAsync(eq(tenantId), any());
+                doReturn(dashboard).when(dashboardService).findDashboardById(eq(tenantId), any());
 
                 break;
             case EDGE:
@@ -220,7 +220,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 edge.setTenantId(tenantId);
 
                 when(ctx.getEdgeService()).thenReturn(edgeService);
-                doReturn(Futures.immediateFuture(edge)).when(edgeService).findEdgeByIdAsync(eq(tenantId), any());
+                doReturn(edge).when(edgeService).findEdgeById(eq(tenantId), any());
 
                 break;
             case OTA_PACKAGE:
@@ -228,7 +228,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 otaPackage.setTenantId(tenantId);
 
                 when(ctx.getOtaPackageService()).thenReturn(otaPackageService);
-                doReturn(Futures.immediateFuture(otaPackage)).when(otaPackageService).findOtaPackageInfoByIdAsync(eq(tenantId), any());
+                doReturn(otaPackage).when(otaPackageService).findOtaPackageInfoById(eq(tenantId), any());
 
                 break;
             case ASSET_PROFILE:
@@ -264,11 +264,11 @@ public class EntitiesTenantIdAsyncLoaderTest {
 
                 break;
             case RPC:
-                Rpc rps = new Rpc();
-                rps.setTenantId(tenantId);
+                Rpc rpc = new Rpc();
+                rpc.setTenantId(tenantId);
 
                 when(ctx.getRpcService()).thenReturn(rpcService);
-                doReturn(Futures.immediateFuture(rps)).when(rpcService).findRpcByIdAsync(eq(tenantId), any());
+                doReturn(rpc).when(rpcService).findRpcById(eq(tenantId), any());
 
                 break;
             case QUEUE:
@@ -292,7 +292,7 @@ public class EntitiesTenantIdAsyncLoaderTest {
                 tbResource.setTenantId(tenantId);
 
                 when(ctx.getResourceService()).thenReturn(resourceService);
-                doReturn(Futures.immediateFuture(tbResource)).when(resourceService).findResourceInfoByIdAsync(eq(tenantId), any());
+                doReturn(tbResource).when(resourceService).findResourceInfoById(eq(tenantId), any());
 
                 break;
             case RULE_NODE:
@@ -303,10 +303,9 @@ public class EntitiesTenantIdAsyncLoaderTest {
 
                 break;
             case TENANT_PROFILE:
-                TenantProfile tenantProfile = new TenantProfile();
+                TenantProfile tenantProfile = new TenantProfile(tenantProfileId);
 
                 when(ctx.getTenantProfile()).thenReturn(tenantProfile);
-                when(ctx.getTenantProfile(any(TenantId.class))).thenReturn(tenantProfile);
 
                 break;
             default:
@@ -319,12 +318,17 @@ public class EntitiesTenantIdAsyncLoaderTest {
         return EntityIdFactory.getByTypeAndUuid(entityType, UUID.randomUUID());
     }
 
-    private void checkTenant(TenantId checkTenantId, boolean equals) throws ExecutionException, InterruptedException {
+    private void checkTenant(TenantId checkTenantId, boolean equals) {
         for (EntityType entityType : EntityType.values()) {
-            EntityId entityId = EntityType.TENANT.equals(entityType) ? tenantId : getEntityId(entityType);
-            TenantId targetTenantId = EntitiesTenantIdAsyncLoader.findEntityIdAsync(ctx, entityId).get();
-
-            Assert.assertNotNull(targetTenantId);
+            EntityId entityId;
+            if(EntityType.TENANT.equals(entityType)){
+                entityId = tenantId;
+            } else if(EntityType.TENANT_PROFILE.equals(entityType)){
+                entityId = tenantProfileId;
+            } else {
+                entityId = getEntityId(entityType);
+            }
+            TenantId targetTenantId = TenantIdLoader.findTenantId(ctx, entityId);
             String msg = "Check entity type <" + entityType.name() + ">:";
             if (equals) {
                 Assert.assertEquals(msg, targetTenantId, checkTenantId);
@@ -335,12 +339,12 @@ public class EntitiesTenantIdAsyncLoaderTest {
     }
 
     @Test
-    public void test_findEntityIdAsync_current_tenant() throws ExecutionException, InterruptedException {
+    public void test_findEntityIdAsync_current_tenant() {
         checkTenant(tenantId, true);
     }
 
     @Test
-    public void test_findEntityIdAsync_other_tenant() throws ExecutionException, InterruptedException {
+    public void test_findEntityIdAsync_other_tenant() {
         checkTenant(new TenantId(UUID.randomUUID()), false);
     }