Search in sources :

Example 16 with TargetState

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();
}
Also used : TargetState(org.apache.kafka.connect.runtime.TargetState) FutureCallback(org.apache.kafka.connect.util.FutureCallback) Callback(org.apache.kafka.connect.util.Callback) ConnectorTaskId(org.apache.kafka.connect.util.ConnectorTaskId) RestartPlan(org.apache.kafka.connect.runtime.RestartPlan) RestartRequest(org.apache.kafka.connect.runtime.RestartRequest) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Aggregations

TargetState (org.apache.kafka.connect.runtime.TargetState)16 HerderConnectorContext (org.apache.kafka.connect.runtime.HerderConnectorContext)7 ConnectorTaskId (org.apache.kafka.connect.util.ConnectorTaskId)7 HashMap (java.util.HashMap)4 Map (java.util.Map)4 LogContext (org.apache.kafka.common.utils.LogContext)4 RestartPlan (org.apache.kafka.connect.runtime.RestartPlan)4 RestartRequest (org.apache.kafka.connect.runtime.RestartRequest)4 Callback (org.apache.kafka.connect.util.Callback)4 ConnectException (org.apache.kafka.connect.errors.ConnectException)3 NotFoundException (org.apache.kafka.connect.errors.NotFoundException)3 ArrayList (java.util.ArrayList)2 Arrays (java.util.Arrays)2 Collection (java.util.Collection)2 Collections (java.util.Collections)2 HashSet (java.util.HashSet)2 List (java.util.List)2 NavigableSet (java.util.NavigableSet)2 NoSuchElementException (java.util.NoSuchElementException)2 Objects (java.util.Objects)2