From e7e119b55920307f03a7d0e05c1864daa75cf080 Mon Sep 17 00:00:00 2001 From: Zoltan Haindrich Date: Thu, 16 May 2024 13:33:27 +0000 Subject: [PATCH] reduce copypaste --- .../apache/druid/quidem/DiscovertModule.java | 89 +------------------ .../druid/sql/calcite/util/CalciteTests.java | 24 +++-- 2 files changed, 21 insertions(+), 92 deletions(-) diff --git a/integration-tests/src/main/java/org/apache/druid/quidem/DiscovertModule.java b/integration-tests/src/main/java/org/apache/druid/quidem/DiscovertModule.java index fea9b22265b..94bb2b1d2a5 100644 --- a/integration-tests/src/main/java/org/apache/druid/quidem/DiscovertModule.java +++ b/integration-tests/src/main/java/org/apache/druid/quidem/DiscovertModule.java @@ -20,14 +20,9 @@ package org.apache.druid.quidem; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Sets; import com.google.inject.AbstractModule; import com.google.inject.Provides; -import org.apache.druid.discovery.DiscoveryDruidNode; -import org.apache.druid.discovery.DruidNodeDiscovery; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; -import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.annotations.Json; import org.apache.druid.query.QueryRunnerFactoryConglomerate; @@ -38,12 +33,7 @@ import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.schema.BrokerSegmentMetadataCache; import org.apache.druid.sql.calcite.util.CalciteTests; -import java.util.Collection; -import java.util.HashSet; -import java.util.Map; import java.util.Properties; -import java.util.Set; -import java.util.function.BooleanSupplier; public class DiscovertModule extends AbstractModule { @@ -83,84 +73,11 @@ public class DiscovertModule extends AbstractModule { return new NativeSqlEngine(CalciteTests.createMockQueryLifecycleFactory(walker, conglomerate), jsonMapper); } - @Provides @LazySingleton - DruidNodeDiscoveryProvider getProvider() { - final DruidNode coordinatorNode = new DruidNode("test-coordinator", "dummy", false, 8081, null, true, false); - DiscovertModule.FakeDruidNodeDiscoveryProvider provider = new FakeDruidNodeDiscoveryProvider( - ImmutableMap.of( - NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDINATOR, coordinatorNode)) - ) - ); - return provider; - } - - /** - * A fake {@link DruidNodeDiscoveryProvider} for {@link #createMockSystemSchema}. - */ - private static class FakeDruidNodeDiscoveryProvider extends DruidNodeDiscoveryProvider + DruidNodeDiscoveryProvider getProvider() { - private final Map nodeDiscoveries; - - public FakeDruidNodeDiscoveryProvider(Map nodeDiscoveries) - { - this.nodeDiscoveries = nodeDiscoveries; - } - - @Override - public BooleanSupplier getForNode(DruidNode node, NodeRole nodeRole) - { - boolean get = nodeDiscoveries.getOrDefault(nodeRole, new FakeDruidNodeDiscovery()) - .getAllNodes() - .stream() - .anyMatch(x -> x.getDruidNode().equals(node)); - return () -> get; - } - - @Override - public DruidNodeDiscovery getForNodeRole(NodeRole nodeRole) - { - return nodeDiscoveries.getOrDefault(nodeRole, new FakeDruidNodeDiscovery()); - } + final DruidNode coordinatorNode = CalciteTests.mockCoordinatorNode(); + return CalciteTests.mockDruidNodeDiscoveryProvider(coordinatorNode); } - - private static class FakeDruidNodeDiscovery implements DruidNodeDiscovery - { - private final Set nodes; - - FakeDruidNodeDiscovery() - { - this.nodes = new HashSet<>(); - } - - FakeDruidNodeDiscovery(Map nodes) - { - this.nodes = Sets.newHashSetWithExpectedSize(nodes.size()); - nodes.forEach((k, v) -> { - addNode(v, k); - }); - } - - @Override - public Collection getAllNodes() - { - return nodes; - } - - void addNode(DruidNode node, NodeRole role) - { - final DiscoveryDruidNode discoveryNode = new DiscoveryDruidNode(node, role, ImmutableMap.of()); - this.nodes.add(discoveryNode); - } - - @Override - public void registerListener(Listener listener) - { - - } - } - - - } \ No newline at end of file diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index 2594d2e98dc..ddd853d97be 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -360,18 +360,30 @@ public class CalciteTests return QueryFrameworkUtils.createOperatorTable(INJECTOR); } + + public static DruidNode mockCoordinatorNode() + { + return new DruidNode("test-coordinator", "dummy", false, 8081, null, true, false); + } + + public static FakeDruidNodeDiscoveryProvider mockDruidNodeDiscoveryProvider(final DruidNode coordinatorNode) + { + FakeDruidNodeDiscoveryProvider provider = new FakeDruidNodeDiscoveryProvider( + ImmutableMap.of( + NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDINATOR, coordinatorNode)) + ) + ); + return provider; + } + public static SystemSchema createMockSystemSchema( final DruidSchema druidSchema, final SpecificSegmentsQuerySegmentWalker walker, final AuthorizerMapper authorizerMapper ) { - final DruidNode coordinatorNode = new DruidNode("test-coordinator", "dummy", false, 8081, null, true, false); - FakeDruidNodeDiscoveryProvider provider = new FakeDruidNodeDiscoveryProvider( - ImmutableMap.of( - NodeRole.COORDINATOR, new FakeDruidNodeDiscovery(ImmutableMap.of(NodeRole.COORDINATOR, coordinatorNode)) - ) - ); + final DruidNode coordinatorNode = mockCoordinatorNode(); + FakeDruidNodeDiscoveryProvider provider = mockDruidNodeDiscoveryProvider(coordinatorNode); final DruidNode overlordNode = new DruidNode("test-overlord", "dummy", false, 8090, null, true, false);