reduce copypaste

kgyrtkirk · May 16, 2024 · e7e119b · e7e119b
1 parent fc9a6c7
commit e7e119b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 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
@@ -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.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 SqlEngine createMockSqlEngine(
       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
-    {
-      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());
-      }
-    }
-
-    private static class FakeDruidNodeDiscovery implements DruidNodeDiscovery
+    DruidNodeDiscoveryProvider getProvider()
     {
-      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)
-      {
-
-      }
+      final DruidNode coordinatorNode = CalciteTests.mockCoordinatorNode();
+      return CalciteTests.mockDruidNodeDiscoveryProvider(coordinatorNode);
     }
-
-
-
   }
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
@@ -360,18 +360,30 @@ public static DruidOperatorTable createOperatorTable()
     return QueryFrameworkUtils.createOperatorTable(INJECTOR);
   }
 
-  public static SystemSchema createMockSystemSchema(
-      final DruidSchema druidSchema,
-      final SpecificSegmentsQuerySegmentWalker walker,
-      final AuthorizerMapper authorizerMapper
-  )
+
+  public static DruidNode mockCoordinatorNode()
+  {
+    return new DruidNode("test-coordinator", "dummy", false, 8081, null, true, false);
+  }
+
+  public static FakeDruidNodeDiscoveryProvider mockDruidNodeDiscoveryProvider(final DruidNode coordinatorNode)
   {
-    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))
         )
     );
+    return provider;
+  }
+
+  public static SystemSchema createMockSystemSchema(
+      final DruidSchema druidSchema,
+      final SpecificSegmentsQuerySegmentWalker walker,
+      final AuthorizerMapper authorizerMapper
+  )
+  {
+    final DruidNode coordinatorNode = mockCoordinatorNode();
+    FakeDruidNodeDiscoveryProvider provider = mockDruidNodeDiscoveryProvider(coordinatorNode);
 
     final DruidNode overlordNode = new DruidNode("test-overlord", "dummy", false, 8090, null, true, false);