diff --git a/application/src/test/java/org/thingsboard/server/queue/discovery/HashPartitionServiceTest.java b/application/src/test/java/org/thingsboard/server/queue/discovery/HashPartitionServiceTest.java index e67bbb9506..4d389b4210 100644 --- a/application/src/test/java/org/thingsboard/server/queue/discovery/HashPartitionServiceTest.java +++ b/application/src/test/java/org/thingsboard/server/queue/discovery/HashPartitionServiceTest.java @@ -63,7 +63,6 @@ import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -73,9 +72,9 @@ public class HashPartitionServiceTest { public static final int ITERATIONS = 1000000; public static final int SERVER_COUNT = 3; - private HashPartitionService clusterRoutingService; + private HashPartitionService partitionService; - private TbServiceInfoProvider discoveryService; + private TbServiceInfoProvider serviceInfoProvider; private TenantRoutingInfoService routingInfoService; private ApplicationEventPublisher applicationEventPublisher; private QueueRoutingInfoService queueRoutingInfoService; @@ -85,19 +84,17 @@ public class HashPartitionServiceTest { @Before public void setup() throws Exception { - discoveryService = mock(TbServiceInfoProvider.class); + serviceInfoProvider = mock(TbServiceInfoProvider.class); applicationEventPublisher = mock(ApplicationEventPublisher.class); routingInfoService = mock(TenantRoutingInfoService.class); queueRoutingInfoService = mock(QueueRoutingInfoService.class); topicService = mock(TopicService.class); when(topicService.buildTopicName(Mockito.any())).thenAnswer(i -> i.getArguments()[0]); - clusterRoutingService = createPartitionService(); + partitionService = createPartitionService(); ServiceInfo currentServer = ServiceInfo.newBuilder() .setServiceId("tb-core-0") .addAllServiceTypes(Collections.singletonList(ServiceType.TB_CORE.name())) .build(); -// when(queueService.resolve(Mockito.any(), Mockito.anyString())).thenAnswer(i -> i.getArguments()[1]); -// when(discoveryService.getServiceInfo()).thenReturn(currentServer); List otherServers = new ArrayList<>(); for (int i = 1; i < SERVER_COUNT; i++) { otherServers.add(ServiceInfo.newBuilder() @@ -106,7 +103,7 @@ public class HashPartitionServiceTest { .build()); } - clusterRoutingService.recalculatePartitions(currentServer, otherServers); + partitionService.recalculatePartitions(currentServer, otherServers); } @Test @@ -122,7 +119,7 @@ public class HashPartitionServiceTest { long start = System.currentTimeMillis(); Map map = new HashMap<>(); for (DeviceId deviceId : devices) { - TopicPartitionInfo address = clusterRoutingService.resolve(ServiceType.TB_CORE, TenantId.SYS_TENANT_ID, deviceId); + TopicPartitionInfo address = partitionService.resolve(ServiceType.TB_CORE, TenantId.SYS_TENANT_ID, deviceId); Integer partition = address.getPartition().get(); map.put(partition, map.getOrDefault(partition, 0) + 1); } @@ -156,7 +153,7 @@ public class HashPartitionServiceTest { for (int queueIndex = 0; queueIndex < queueCount; queueIndex++) { QueueKey queueKey = new QueueKey(ServiceType.TB_RULE_ENGINE, "queue" + queueIndex, tenantId); for (int partition = 0; partition < partitionCount; partition++) { - ServiceInfo serviceInfo = clusterRoutingService.resolveByPartitionIdx(services, queueKey, partition, Collections.emptyMap()); + ServiceInfo serviceInfo = partitionService.resolveByPartitionIdx(services, queueKey, partition, Collections.emptyMap()); String serviceId = serviceInfo.getServiceId(); map.put(serviceId, map.get(serviceId) + 1); } @@ -233,15 +230,34 @@ public class HashPartitionServiceTest { } Map>> serversPartitions = new HashMap<>(); - clusterRoutingService.init(); + when(serviceInfoProvider.isService(eq(ServiceType.TB_RULE_ENGINE))).thenReturn(true); + partitionService.init(); for (ServiceInfo ruleEngine : ruleEngines) { List other = new ArrayList<>(ruleEngines); other.removeIf(serviceInfo -> serviceInfo.getServiceId().equals(ruleEngine.getServiceId())); - clusterRoutingService.recalculatePartitions(ruleEngine, other); - clusterRoutingService.myPartitions.forEach((queueKey, partitions) -> { + partitionService.recalculatePartitions(ruleEngine, other); + partitionService.myPartitions.forEach((queueKey, partitions) -> { serversPartitions.computeIfAbsent(queueKey, k -> new HashMap<>()).put(ruleEngine, partitions); }); + + Set assignedTenantProfiles = ruleEngine.getAssignedTenantProfilesList().stream().map(UUID::fromString).collect(Collectors.toSet()); + when(serviceInfoProvider.getAssignedTenantProfiles()).thenReturn(assignedTenantProfiles); + if (assignedTenantProfiles.isEmpty()) { + assertThat(partitionService.isManagedByCurrentService(TenantId.SYS_TENANT_ID)).isTrue(); + tenants.forEach((tenantId, tenantProfileId) -> { + assertThat(partitionService.isManagedByCurrentService(tenantId)).isFalse(); + }); + } else { + assertThat(partitionService.isManagedByCurrentService(TenantId.SYS_TENANT_ID)).isFalse(); + tenants.forEach((tenantId, tenantProfileId) -> { + if (assignedTenantProfiles.contains(tenantProfileId.getId())) { + assertThat(partitionService.isManagedByCurrentService(tenantId)).isTrue(); + } else { + assertThat(partitionService.isManagedByCurrentService(tenantId)).isFalse(); + } + }); + } } assertThat(serversPartitions.keySet()).containsAll(queues.stream().map(queue -> new QueueKey(ServiceType.TB_RULE_ENGINE, queue)).collect(Collectors.toList())); @@ -286,7 +302,7 @@ public class HashPartitionServiceTest { mockRoutingInfo(tenantId, tenantProfileId, false); // not isolated yet mockQueues(queues); - when(discoveryService.isService(eq(ServiceType.TB_RULE_ENGINE))).thenReturn(true); + when(serviceInfoProvider.isService(eq(ServiceType.TB_RULE_ENGINE))).thenReturn(true); Mockito.reset(applicationEventPublisher); HashPartitionService partitionService_common = createPartitionService(); partitionService_common.recalculatePartitions(commonRuleEngine, List.of(dedicatedRuleEngine)); @@ -349,27 +365,6 @@ public class HashPartitionServiceTest { }); } - @Test - public void testIsManagedByCurrentServiceCheck() { - TenantProfileId isolatedProfileId = new TenantProfileId(UUID.randomUUID()); - when(discoveryService.getAssignedTenantProfiles()).thenReturn(Set.of(isolatedProfileId.getId())); // dedicated server - TenantProfileId regularProfileId = new TenantProfileId(UUID.randomUUID()); - - TenantId isolatedTenantId = new TenantId(UUID.randomUUID()); - mockRoutingInfo(isolatedTenantId, isolatedProfileId, true); - TenantId regularTenantId = new TenantId(UUID.randomUUID()); - mockRoutingInfo(regularTenantId, regularProfileId, false); - - assertThat(clusterRoutingService.isManagedByCurrentService(isolatedTenantId)).isTrue(); - assertThat(clusterRoutingService.isManagedByCurrentService(regularTenantId)).isFalse(); - - - when(discoveryService.getAssignedTenantProfiles()).thenReturn(Collections.emptySet()); // common server - - assertThat(clusterRoutingService.isManagedByCurrentService(isolatedTenantId)).isTrue(); - assertThat(clusterRoutingService.isManagedByCurrentService(regularTenantId)).isTrue(); - } - @Test public void testPartitionsDistribution_sameTenantDifferentQueues() { List ruleEngines = new ArrayList<>(); @@ -389,9 +384,9 @@ public class HashPartitionServiceTest { .limit(100).collect(Collectors.toList()); for (int partition = 0; partition < 10; partition++) { - ServiceInfo expectedAssignedRuleEngine = clusterRoutingService.resolveByPartitionIdx(ruleEngines, new QueueKey(ServiceType.TB_RULE_ENGINE, tenantId), partition, Collections.emptyMap()); + ServiceInfo expectedAssignedRuleEngine = partitionService.resolveByPartitionIdx(ruleEngines, new QueueKey(ServiceType.TB_RULE_ENGINE, tenantId), partition, Collections.emptyMap()); for (QueueKey queueKey : queues) { - ServiceInfo assignedRuleEngine = clusterRoutingService.resolveByPartitionIdx(ruleEngines, queueKey, partition, Collections.emptyMap()); + ServiceInfo assignedRuleEngine = partitionService.resolveByPartitionIdx(ruleEngines, queueKey, partition, Collections.emptyMap()); assertThat(assignedRuleEngine).as(queueKey + "[" + partition + "] should be assigned to " + expectedAssignedRuleEngine.getServiceId()) .isEqualTo(expectedAssignedRuleEngine); } @@ -403,9 +398,9 @@ public class HashPartitionServiceTest { verify(applicationEventPublisher).publishEvent(argThat(event -> event instanceof PartitionChangeEvent && predicate.test((PartitionChangeEvent) event))); } - private void mockRoutingInfo(TenantId tenantId, TenantProfileId tenantProfileId, boolean isolatedTbRuleEngine) { + private void mockRoutingInfo(TenantId tenantId, TenantProfileId tenantProfileId, boolean isolated) { when(routingInfoService.getRoutingInfo(eq(tenantId))) - .thenReturn(new TenantRoutingInfo(tenantId, tenantProfileId, isolatedTbRuleEngine)); + .thenReturn(new TenantRoutingInfo(tenantId, tenantProfileId, isolated)); } private void mockQueues(List queues) { @@ -424,7 +419,7 @@ public class HashPartitionServiceTest { } private HashPartitionService createPartitionService() { - HashPartitionService partitionService = new HashPartitionService(discoveryService, + HashPartitionService partitionService = new HashPartitionService(serviceInfoProvider, routingInfoService, applicationEventPublisher, queueRoutingInfoService,