use of org.apache.kafka.connect.runtime.TargetState in project kafka by apache.
the class DistributedHerderTest method testDoRestartConnectorAndTasksBoth.
@Test
public void testDoRestartConnectorAndTasksBoth() {
ConnectorTaskId taskId = new ConnectorTaskId(CONN1, 0);
RestartRequest restartRequest = new RestartRequest(CONN1, false, true);
RestartPlan restartPlan = PowerMock.createMock(RestartPlan.class);
EasyMock.expect(restartPlan.shouldRestartConnector()).andReturn(true).anyTimes();
EasyMock.expect(restartPlan.shouldRestartTasks()).andReturn(true).anyTimes();
EasyMock.expect(restartPlan.taskIdsToRestart()).andReturn(Collections.singletonList(taskId)).anyTimes();
EasyMock.expect(restartPlan.restartTaskCount()).andReturn(1).anyTimes();
EasyMock.expect(restartPlan.totalTaskCount()).andReturn(1).anyTimes();
EasyMock.expect(herder.buildRestartPlan(restartRequest)).andReturn(Optional.of(restartPlan)).anyTimes();
herder.assignment = PowerMock.createMock(ExtendedAssignment.class);
EasyMock.expect(herder.assignment.connectors()).andReturn(Collections.singletonList(CONN1)).anyTimes();
EasyMock.expect(herder.assignment.tasks()).andReturn(Collections.singletonList(taskId)).anyTimes();
worker.stopAndAwaitConnector(CONN1);
PowerMock.expectLastCall();
Capture<Callback<TargetState>> stateCallback = newCapture();
worker.startConnector(EasyMock.eq(CONN1), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.eq(herder), EasyMock.anyObject(TargetState.class), capture(stateCallback));
herder.onRestart(CONN1);
EasyMock.expectLastCall();
worker.stopAndAwaitTasks(Collections.singletonList(taskId));
PowerMock.expectLastCall();
herder.onRestart(taskId);
EasyMock.expectLastCall();
worker.startTask(EasyMock.eq(TASK0), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.anyObject(), EasyMock.eq(herder), EasyMock.anyObject(TargetState.class));
PowerMock.expectLastCall().andReturn(true);
PowerMock.replayAll();
herder.doRestartConnectorAndTasks(restartRequest);
PowerMock.verifyAll();
}
Aggregations