Factor out common code in TestUniformNodeSelector

trinodb · Dec 21, 2023 · 2b30080 · 2b30080
1 parent 6afc644
commit 2b30080
Showing 1 changed file with 5 additions and 21 deletions.
diff --git a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestUniformNodeSelector.java b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestUniformNodeSelector.java
@@ -64,6 +64,9 @@
 @TestInstance(PER_METHOD)
 public class TestUniformNodeSelector
 {
+    private static final InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:13"), NodeVersion.UNKNOWN, false);
+    private static final InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
+    private final Set<Split> splits = new LinkedHashSet<>();
     private FinalizerService finalizerService;
     private NodeTaskMap nodeTaskMap;
     private InMemoryNodeManager nodeManager;
@@ -82,6 +85,8 @@ public void setUp()
         finalizerService = new FinalizerService();
         nodeTaskMap = new NodeTaskMap(finalizerService);
         nodeManager = new InMemoryNodeManager();
+        nodeManager.addNodes(node1);
+        nodeManager.addNodes(node2);
 
         nodeSchedulerConfig = new NodeSchedulerConfig()
                 .setMaxSplitsPerNode(20)
@@ -116,11 +121,6 @@ public void tearDown()
     @Test
     public void testQueueSizeAdjustmentScaleDown()
     {
-        InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:13"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node1);
-        InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node2);
-
         TestingTicker ticker = new TestingTicker();
         UniformNodeSelector.QueueSizeAdjuster queueSizeAdjuster = new UniformNodeSelector.QueueSizeAdjuster(10, 100, ticker);
 
@@ -137,8 +137,6 @@ public void testQueueSizeAdjustmentScaleDown()
                 false,
                 queueSizeAdjuster);
 
-        Set<Split> splits = new LinkedHashSet<>();
-
         for (int i = 0; i < 20; i++) {
             splits.add(new Split(TEST_CATALOG_HANDLE, TestingSplit.createRemoteSplit()));
         }
@@ -187,13 +185,6 @@ public void testQueueSizeAdjustmentScaleDown()
     @Test
     public void testQueueSizeAdjustmentAllNodes()
     {
-        InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:13"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node1);
-        InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node2);
-
-        Set<Split> splits = new LinkedHashSet<>();
-
         for (int i = 0; i < 20 * 9; i++) {
             splits.add(new Split(TEST_CATALOG_HANDLE, TestingSplit.createRemoteSplit()));
         }
@@ -246,13 +237,6 @@ public void testQueueSizeAdjustmentAllNodes()
     @Test
     public void testQueueSizeAdjustmentOneOfAll()
     {
-        InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:13"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node1);
-        InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
-        nodeManager.addNodes(node2);
-
-        Set<Split> splits = new LinkedHashSet<>();
-
         for (int i = 0; i < 20 * 9; i++) {
             splits.add(new Split(TEST_CATALOG_HANDLE, TestingSplit.createRemoteSplit()));
         }