@Test
public void testUpdateShuffleNode() throws Exception {
JobId jobId = JobId.newJobId(createTestAppId(), 123456);
DragonZooKeeper.NodeData node = new DragonZooKeeper.NodeData();
node.nodeId = createTestNodeId();
node.taskId = TaskId.newTaskId(jobId, 34565, TaskType.MAP);
dragonZK.createShufflePath(jobId);
dragonZK.createShuffleNode(jobId, newArrayList(node));
node.nodeId.setHost("NewServerHost");
dragonZK.updateShuffleNode(jobId, newArrayList(node));
Stat stat =
zkClient
.checkExists()
.forPath(
"/dragon/job_458495849584_89894985_123456/shuffle/task_458495849584_89894985_123456_m_34565");
assertNotNull(stat);
byte[] data =
zkClient
.getData()
.forPath(
"/dragon/job_458495849584_89894985_123456/shuffle/task_458495849584_89894985_123456_m_34565");
assertEquals("NewServerHost:2345", new String(data));
}
@Test
public void testGetAvailableNodes() throws Exception {
NodeId nodeId = createTestNodeId();
dragonZK.registerNodeManager(nodeId);
List<NodeId> nodeIdList = dragonZK.getAvailableNodes();
assertEquals(nodeId, nodeIdList.get(0));
}
@Test
public void testRegisterNodeManager() throws Exception {
NodeId nodeId = createTestNodeId();
dragonZK.registerNodeManager(nodeId);
Stat stat = zkClient.checkExists().forPath("/dragon/nodemanagers/ServerHost:2345");
assertNotNull(stat);
}
@Test
public void testGetShufflePath() throws Exception {
ApplicationId appId = createTestAppId();
JobId jobId = JobId.newJobId(appId, 123456);
assertEquals(
"/dragon/job_458495849584_89894985_123456/shuffle", dragonZK.getShufflePath(jobId));
}
@Test
public void testWatchNodeManager() throws Exception {
NodeId nodeId = createTestNodeId();
dragonZK.registerNodeManager(nodeId);
nodeDied = false;
dragonZK.watchNodeManager(
nodeId,
new CuratorWatcher() {
@Override
public void process(WatchedEvent event) throws Exception {
nodeDied = true;
}
});
zkClient.delete().forPath("/dragon/nodemanagers/ServerHost:2345");
Thread.sleep(1000);
assertTrue(nodeDied);
}
@Test
public void testGetTaskPath() throws Exception {
ApplicationId appId = createTestAppId();
JobId jobId = JobId.newJobId(appId, 123456);
TaskId taskId = TaskId.newTaskId(jobId, 34565, TaskType.MAP);
assertEquals(
"/dragon/job_458495849584_89894985_123456/shuffle/task_458495849584_89894985_123456_m_34565",
dragonZK.getTaskPath(jobId, taskId));
}
@Test
public void testCreateShufflePath() throws Exception {
ApplicationId appId = createTestAppId();
JobId jobId = JobId.newJobId(appId, 123456);
dragonZK.createShufflePath(jobId);
Stat stat = zkClient.checkExists().forPath("/dragon/job_458495849584_89894985_123456/shuffle");
assertNotNull(stat);
}
@Test
public void testGetShuffleNodeByTaskId() throws Exception {
JobId jobId = JobId.newJobId(createTestAppId(), 123456);
TaskId taskId = TaskId.newTaskId(jobId, 34565, TaskType.MAP);
DragonZooKeeper.NodeData node = new DragonZooKeeper.NodeData();
node.nodeId = createTestNodeId();
node.taskId = taskId;
if (zkClient.checkExists().forPath("/dragon/job_458495849584_89894985_123456/shuffle")
== null) {
dragonZK.createShufflePath(jobId);
}
PathChildrenCache shuffleNodeCache =
new PathChildrenCache(zkClient, dragonZK.getShufflePath(jobId), true);
dragonZK.setShuffleNodeCache(shuffleNodeCache);
dragonZK.createShuffleNode(jobId, newArrayList(node));
Thread.sleep(1000);
NodeId nodeId = dragonZK.getShuffleNodeByTaskId(jobId, taskId);
assertNotNull(nodeId);
assertEquals(node.nodeId, nodeId);
}
@Test
public void testGetNodeManagerPath() throws Exception {
NodeId nodeId = createTestNodeId();
assertEquals("/dragon/nodemanagers/ServerHost:2345", dragonZK.getNodeManagerPath(nodeId));
}
@After
public void tearDown() throws Exception {
dragonZK.close();
zkServer.stop();
}