diff --git a/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/util/EntitiesFieldsAsyncLoaderTest.java b/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/util/EntitiesFieldsAsyncLoaderTest.java index 5967023354..f3e61dcffc 100644 --- a/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/util/EntitiesFieldsAsyncLoaderTest.java +++ b/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/util/EntitiesFieldsAsyncLoaderTest.java @@ -37,10 +37,12 @@ import org.thingsboard.server.common.data.Tenant; 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.edge.Edge; import org.thingsboard.server.common.data.id.AlarmId; import org.thingsboard.server.common.data.id.AssetId; import org.thingsboard.server.common.data.id.CustomerId; import org.thingsboard.server.common.data.id.DeviceId; +import org.thingsboard.server.common.data.id.EdgeId; import org.thingsboard.server.common.data.id.EntityId; import org.thingsboard.server.common.data.id.EntityIdFactory; import org.thingsboard.server.common.data.id.EntityViewId; @@ -52,6 +54,7 @@ import org.thingsboard.server.common.data.rule.RuleChain; import org.thingsboard.server.dao.asset.AssetService; import org.thingsboard.server.dao.customer.CustomerService; 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.rule.RuleChainService; import org.thingsboard.server.dao.tenant.TenantService; @@ -95,6 +98,8 @@ public class EntitiesFieldsAsyncLoaderTest { private RuleChainService ruleChainServiceMock; @Mock private EntityViewService entityViewServiceMock; + @Mock + private EdgeService edgeServiceMock; @BeforeAll public static void setup() { @@ -108,7 +113,8 @@ public class EntitiesFieldsAsyncLoaderTest { EntityType.DEVICE, EntityType.ALARM, EntityType.RULE_CHAIN, - EntityType.ENTITY_VIEW + EntityType.ENTITY_VIEW, + EntityType.EDGE ); } @@ -228,6 +234,14 @@ public class EntitiesFieldsAsyncLoaderTest { when(ctxMock.getEntityViewService()).thenReturn(entityViewServiceMock); doReturn(entityView).when(entityViewServiceMock).findEntityViewByIdAsync(eq(TENANT_ID), any()); + break; + case EDGE: + var edge = Futures.immediateFuture(entityDoesNotExist ? null : new Edge(new EdgeId(RANDOM_UUID))); + + when(ctxMock.getDbCallbackExecutor()).thenReturn(DB_EXECUTOR); + when(ctxMock.getEdgeService()).thenReturn(edgeServiceMock); + doReturn(edge).when(edgeServiceMock).findEdgeByIdAsync(eq(TENANT_ID), any()); + break; default: throw new RuntimeException("Unexpected EntityType: " + entityType); @@ -252,6 +266,8 @@ public class EntitiesFieldsAsyncLoaderTest { return new RuleChain((RuleChainId) entityId); case ENTITY_VIEW: return new EntityView((EntityViewId) entityId); + case EDGE: + return new Edge((EdgeId) entityId); default: throw new RuntimeException("Unexpected EntityType: " + entityId.getEntityType()); }