reduce copypaste

This commit is contained in:
Zoltan Haindrich 2024-05-16 13:33:27 +00:00
parent fc9a6c7740
commit e7e119b559
2 changed files with 21 additions and 92 deletions

View File

@ -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<NodeRole, DiscovertModule.FakeDruidNodeDiscovery> nodeDiscoveries;
public FakeDruidNodeDiscoveryProvider(Map<NodeRole, DiscovertModule.FakeDruidNodeDiscovery> 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<DiscoveryDruidNode> nodes;
FakeDruidNodeDiscovery()
{
this.nodes = new HashSet<>();
}
FakeDruidNodeDiscovery(Map<NodeRole, DruidNode> nodes)
{
this.nodes = Sets.newHashSetWithExpectedSize(nodes.size());
nodes.forEach((k, v) -> {
addNode(v, k);
});
}
@Override
public Collection<DiscoveryDruidNode> 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)
{
}
}
}

View File

@ -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);