@Test
public void testSplitAssignment()
throws Exception
{
final InMemoryNodeManager nodeManager = new InMemoryNodeManager();
ImmutableList.Builder<Node> nodeBuilder = ImmutableList.builder();
nodeBuilder.add(new PrestoNode("other1", URI.create("http://127.0.0.1:11"), NodeVersion.UNKNOWN));
nodeBuilder.add(new PrestoNode("other2", URI.create("http://127.0.0.1:12"), NodeVersion.UNKNOWN));
nodeBuilder.add(new PrestoNode("other3", URI.create("http://127.0.0.1:13"), NodeVersion.UNKNOWN));
ImmutableList<Node> nodes = nodeBuilder.build();
nodeManager.addNode("foo", nodes);
NodeScheduler nodeScheduler = new NodeScheduler(nodeManager, new NodeSchedulerConfig());
// Start sql stage execution
SqlStageExecution sqlStageExecution1 = createSqlStageExecution(nodeScheduler, 2, 20);
Future future1 = sqlStageExecution1.start();
future1.get(1, TimeUnit.SECONDS);
Map<Node, RemoteTask> tasks1 = sqlStageExecution1.getTasks();
for (Map.Entry<Node, RemoteTask> entry : tasks1.entrySet()) {
assertEquals(entry.getValue().getQueuedSplits(), 5);
}
// Add new node
Node additionalNode = new PrestoNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN);
nodeManager.addNode("foo", additionalNode);
// Schedule next query
SqlStageExecution sqlStageExecution2 = createSqlStageExecution(nodeScheduler, 2, 20);
Future future2 = sqlStageExecution2.start();
future2.get(1, TimeUnit.SECONDS);