Examples with EdgeManagerForTest org.apache.tez.test.EdgeManagerForTest used on opensource projects

Example 1 with EdgeManagerForTest

use of org.apache.tez.test.EdgeManagerForTest in project tez by apache.

the class TestVertexImpl method testVertexSetParallelismIncrease.

@Test(timeout = 5000)
public void testVertexSetParallelismIncrease() throws Exception {
    VertexImpl v3 = vertices.get("vertex3");
    v3.vertexReconfigurationPlanned();
    initAllVertices(VertexState.INITED);
    Assert.assertEquals(2, v3.getTotalTasks());
    Assert.assertEquals(2, v3.getTasks().size());
    VertexImpl v1 = vertices.get("vertex1");
    startVertex(vertices.get("vertex2"));
    startVertex(v1);
    EdgeManagerPluginDescriptor mockEdgeManagerDescriptor = EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    EdgeProperty edgeProp = EdgeProperty.create(mockEdgeManagerDescriptor, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"), InputDescriptor.create("In"));
    Map<String, EdgeProperty> edgeManagerDescriptors = Collections.singletonMap(v1.getName(), edgeProp);
    v3.reconfigureVertex(10, null, edgeManagerDescriptors);
    v3.doneReconfiguringVertex();
    assertTrue(v3.sourceVertices.get(v1).getEdgeManager() instanceof EdgeManagerForTest);
    checkTasks(v3, 10);
}

Example 2 with EdgeManagerForTest

use of org.apache.tez.test.EdgeManagerForTest in project tez by apache.

the class TestVertexImpl method testSetCustomEdgeManager.

@Test(timeout = 5000)
public void testSetCustomEdgeManager() throws Exception {
    // Vertex5 linked to v3 (v3 src, v5 dest)
    VertexImpl v5 = vertices.get("vertex5");
    v5.vertexReconfigurationPlanned();
    initAllVertices(VertexState.INITED);
    Edge edge = edges.get("e4");
    EdgeManagerPlugin em = edge.getEdgeManager();
    EdgeManagerForTest originalEm = (EdgeManagerForTest) em;
    assertTrue(Arrays.equals(edgePayload, originalEm.getEdgeManagerContext().getUserPayload().deepCopyAsArray()));
    UserPayload userPayload = UserPayload.create(ByteBuffer.wrap(new String("foo").getBytes()));
    EdgeManagerPluginDescriptor edgeManagerDescriptor = EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    edgeManagerDescriptor.setUserPayload(userPayload);
    EdgeProperty edgeProp = EdgeProperty.create(edgeManagerDescriptor, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"), InputDescriptor.create("In"));
    Vertex v3 = vertices.get("vertex3");
    Map<String, EdgeProperty> edgeManagerDescriptors = Collections.singletonMap(v3.getName(), edgeProp);
    v5.reconfigureVertex(v5.getTotalTasks() - 1, null, edgeManagerDescriptors);
    v5.doneReconfiguringVertex();
    VertexImpl v5Impl = (VertexImpl) v5;
    EdgeManagerPlugin modifiedEdgeManager = v5Impl.sourceVertices.get(v3).getEdgeManager();
    Assert.assertNotNull(modifiedEdgeManager);
    assertTrue(modifiedEdgeManager instanceof EdgeManagerForTest);
    // Ensure initialize() is called with the correct payload
    assertTrue(Arrays.equals(userPayload.deepCopyAsArray(), ((EdgeManagerForTest) modifiedEdgeManager).getUserPayload().deepCopyAsArray()));
}

Example 3 with EdgeManagerForTest

use of org.apache.tez.test.EdgeManagerForTest in project tez by apache.

the class TestVertexImpl method testVertexSetParallelismDecrease.

@Test(timeout = 5000)
public void testVertexSetParallelismDecrease() throws Exception {
    VertexImpl v3 = vertices.get("vertex3");
    v3.vertexReconfigurationPlanned();
    initAllVertices(VertexState.INITED);
    Assert.assertEquals(2, v3.getTotalTasks());
    Assert.assertEquals(2, v3.getTasks().size());
    VertexImpl v1 = vertices.get("vertex1");
    startVertex(vertices.get("vertex2"));
    startVertex(v1);
    EdgeManagerPluginDescriptor mockEdgeManagerDescriptor = EdgeManagerPluginDescriptor.create(EdgeManagerForTest.class.getName());
    EdgeProperty edgeProp = EdgeProperty.create(mockEdgeManagerDescriptor, DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL, OutputDescriptor.create("Out"), InputDescriptor.create("In"));
    Map<String, EdgeProperty> edgeManagerDescriptors = Collections.singletonMap(v1.getName(), edgeProp);
    v3.reconfigureVertex(1, null, edgeManagerDescriptors);
    v3.doneReconfiguringVertex();
    assertTrue(v3.sourceVertices.get(v1).getEdgeManager() instanceof EdgeManagerForTest);
    checkTasks(v3, 1);
}