diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java index 00d44b045b13d..5ac90d7aade3e 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java @@ -55,17 +55,14 @@ import java.nio.file.Files; import java.time.Duration; import java.util.ArrayList; -import org.easymock.EasyMock; -import org.easymock.EasyMockRunner; -import org.easymock.Mock; -import org.easymock.MockType; import org.hamcrest.Matchers; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; -import org.mockito.Mockito; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; import java.io.File; import java.util.Arrays; @@ -98,15 +95,6 @@ import static org.apache.kafka.streams.processor.internals.TopologyMetadata.UNNAMED_TOPOLOGY; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask; -import static org.easymock.EasyMock.anyObject; -import static org.easymock.EasyMock.anyString; -import static org.easymock.EasyMock.eq; -import static org.easymock.EasyMock.expect; -import static org.easymock.EasyMock.expectLastCall; -import static org.easymock.EasyMock.replay; -import static org.easymock.EasyMock.reset; -import static org.easymock.EasyMock.resetToStrict; -import static org.easymock.EasyMock.verify; import static org.hamcrest.CoreMatchers.hasItem; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.empty; @@ -122,14 +110,20 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import static org.mockito.Mockito.mock; -@RunWith(EasyMockRunner.class) +@RunWith(MockitoJUnitRunner.StrictStubs.class) public class TaskManagerTest { private final String topic1 = "topic1"; @@ -178,21 +172,22 @@ public class TaskManagerTest { final java.util.function.Consumer> noOpResetter = partitions -> { }; - @Mock(type = MockType.STRICT) + @Mock private InternalTopologyBuilder topologyBuilder; - @Mock(type = MockType.DEFAULT) + @Mock private StateDirectory stateDirectory; - @Mock(type = MockType.NICE) + @Mock private ChangelogReader changeLogReader; - @Mock(type = MockType.STRICT) + @Mock private Consumer consumer; - @Mock(type = MockType.STRICT) + @Mock private ActiveTaskCreator activeTaskCreator; - @Mock(type = MockType.NICE) + @Mock private StandbyTaskCreator standbyTaskCreator; - @Mock(type = MockType.NICE) + @Mock private Admin adminClient; - final StateUpdater stateUpdater = Mockito.mock(StateUpdater.class); + @Mock + private StateUpdater stateUpdater; private TaskManager taskManager; private TopologyMetadata topologyMetadata; @@ -228,9 +223,8 @@ private TaskManager setUpTaskManager(final ProcessingMode processingMode, stateUpdaterEnabled ? stateUpdater : null ); taskManager.setMainConsumer(consumer); - reset(topologyBuilder); - expect(topologyBuilder.hasNamedTopology()).andStubReturn(false); - expect(topologyBuilder.nodeToSourceTopics()).andStubReturn(emptyMap()); + when(topologyBuilder.hasNamedTopology()).thenReturn(false); + when(topologyBuilder.nodeToSourceTopics()).thenReturn(emptyMap()); return taskManager; } @@ -246,7 +240,7 @@ public void shouldClassifyExistingTasksWithoutStateUpdater() { taskManager.handleAssignment(activeTasks, standbyTasks); - Mockito.verifyNoInteractions(stateUpdater); + verifyNoInteractions(stateUpdater); } @Test @@ -255,7 +249,7 @@ public void shouldNotUpdateExistingStandbyTaskIfStandbyIsReassignedWithSameInput .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(standbyTask, taskId03Partitions); - Mockito.verify(standbyTask, never()).updateInputPartitions(Mockito.eq(taskId03Partitions), Mockito.any()); + verify(standbyTask, never()).updateInputPartitions(eq(taskId03Partitions), any()); } @Test @@ -264,12 +258,12 @@ public void shouldUpdateExistingStandbyTaskIfStandbyIsReassignedWithDifferentInp .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(standbyTask, taskId04Partitions); - Mockito.verify(standbyTask).updateInputPartitions(Mockito.eq(taskId04Partitions), Mockito.any()); + verify(standbyTask).updateInputPartitions(eq(taskId04Partitions), any()); } private void updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(final Task standbyTask, final Set newInputPartition) { - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); when(tasks.allTasks()).thenReturn(mkSet(standbyTask)); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, false); @@ -278,7 +272,7 @@ private void updateExistingStandbyTaskIfStandbyIsReassignedWithoutStateUpdater(f mkMap(mkEntry(standbyTask.id(), newInputPartition)) ); - Mockito.verify(standbyTask).resume(); + verify(standbyTask).resume(); } @Test @@ -286,20 +280,18 @@ public void shouldPrepareActiveTaskInStateUpdaterToBeRecycled() { final StreamTask activeTaskToRecycle = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToRecycle)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( Collections.emptyMap(), mkMap(mkEntry(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions())) ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(tasks).addPendingTaskToRecycle(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions()); + verify(tasks).addPendingTaskToRecycle(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions()); } @Test @@ -307,21 +299,19 @@ public void shouldPrepareStandbyTaskInStateUpdaterToBeRecycled() { final StandbyTask standbyTaskToRecycle = standbyTask(taskId03, taskId03ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToRecycle)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater).remove(standbyTaskToRecycle.id()); - Mockito.verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions()); + verify(stateUpdater).remove(standbyTaskToRecycle.id()); + verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions()); } @Test @@ -329,18 +319,16 @@ public void shouldRemoveUnusedActiveTaskFromStateUpdater() { final StreamTask activeTaskToClose = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToClose)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap()); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater).remove(activeTaskToClose.id()); - Mockito.verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id()); + verify(stateUpdater).remove(activeTaskToClose.id()); + verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id()); } @Test @@ -348,18 +336,16 @@ public void shouldRemoveUnusedStandbyTaskFromStateUpdater() { final StandbyTask standbyTaskToClose = standbyTask(taskId02, taskId02ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToClose)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap()); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater).remove(standbyTaskToClose.id()); - Mockito.verify(tasks).addPendingTaskToCloseClean(standbyTaskToClose.id()); + verify(stateUpdater).remove(standbyTaskToClose.id()); + verify(tasks).addPendingTaskToCloseClean(standbyTaskToClose.id()); } @Test @@ -368,21 +354,19 @@ public void shouldUpdateInputPartitionOfActiveTaskInStateUpdater() { .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); final Set newInputPartitions = taskId02Partitions; - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToUpdateInputPartitions)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToUpdateInputPartitions.id(), newInputPartitions)), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater).remove(activeTaskToUpdateInputPartitions.id()); - Mockito.verify(tasks).addPendingTaskToUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions); + verify(stateUpdater).remove(activeTaskToUpdateInputPartitions.id()); + verify(tasks).addPendingTaskToUpdateInputPartitions(activeTaskToUpdateInputPartitions.id(), newInputPartitions); } @Test @@ -390,19 +374,16 @@ public void shouldKeepReAssignedActiveTaskInStateUpdater() { final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedActiveTask)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(reassignedActiveTask.id(), reassignedActiveTask.inputPartitions())), Collections.emptyMap() ); - - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -410,20 +391,18 @@ public void shouldRemoveReAssignedRevokedActiveTaskInStateUpdaterFromPendingTask final StreamTask reAssignedRevokedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(reAssignedRevokedActiveTask)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(reAssignedRevokedActiveTask.id(), reAssignedRevokedActiveTask.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(tasks).removePendingActiveTaskToSuspend(reAssignedRevokedActiveTask.id()); + verify(tasks).removePendingActiveTaskToSuspend(reAssignedRevokedActiveTask.id()); } @Test @@ -432,21 +411,19 @@ public void shouldNeverUpdateInputPartitionsOfStandbyTaskInStateUpdater() { .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); final Set newInputPartitions = taskId03Partitions; - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTaskToUpdateInputPartitions)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( Collections.emptyMap(), mkMap(mkEntry(standbyTaskToUpdateInputPartitions.id(), newInputPartitions)) ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id()); - Mockito.verify(tasks, never()) + verify(stateUpdater, never()).remove(standbyTaskToUpdateInputPartitions.id()); + verify(tasks, never()) .addPendingTaskToUpdateInputPartitions(standbyTaskToUpdateInputPartitions.id(), newInputPartitions); } @@ -455,19 +432,16 @@ public void shouldKeepReAssignedStandbyTaskInStateUpdater() { final StandbyTask reAssignedStandbyTask = standbyTask(taskId02, taskId02ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(reAssignedStandbyTask)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( Collections.emptyMap(), mkMap(mkEntry(reAssignedStandbyTask.id(), reAssignedStandbyTask.inputPartitions())) ); - - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -478,23 +452,21 @@ public void shouldAssignMultipleTasksInStateUpdater() { final StandbyTask standbyTaskToRecycle = standbyTask(taskId02, taskId02ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(stateUpdater.getTasks()).thenReturn(mkSet(activeTaskToClose, standbyTaskToRecycle)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(stateUpdater).remove(activeTaskToClose.id()); - Mockito.verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id()); - Mockito.verify(stateUpdater).remove(standbyTaskToRecycle.id()); - Mockito.verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions()); + verify(stateUpdater).remove(activeTaskToClose.id()); + verify(tasks).addPendingTaskToCloseClean(activeTaskToClose.id()); + verify(stateUpdater).remove(standbyTaskToRecycle.id()); + verify(tasks).addPendingTaskToRecycle(standbyTaskToRecycle.id(), standbyTaskToRecycle.inputPartitions()); } @Test @@ -502,22 +474,20 @@ public void shouldCreateActiveTaskDuringAssignment() { final StreamTask activeTaskToBeCreated = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.CREATED) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); final Set createdTasks = mkSet(activeTaskToBeCreated); - expect(activeTaskCreator.createTasks(consumer, mkMap( + when(activeTaskCreator.createTasks(consumer, mkMap( mkEntry(activeTaskToBeCreated.id(), activeTaskToBeCreated.inputPartitions()))) - ).andReturn(createdTasks); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + ).thenReturn(createdTasks); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToBeCreated.id(), activeTaskToBeCreated.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(tasks).addPendingTaskToInit(createdTasks); + verify(tasks).addPendingTaskToInit(createdTasks); } @Test @@ -525,22 +495,20 @@ public void shouldCreateStandbyTaskDuringAssignment() { final StandbyTask standbyTaskToBeCreated = standbyTask(taskId02, taskId02ChangelogPartitions) .inState(State.CREATED) .withInputPartitions(taskId02Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); final Set createdTasks = mkSet(standbyTaskToBeCreated); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(mkMap( + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(mkMap( mkEntry(standbyTaskToBeCreated.id(), standbyTaskToBeCreated.inputPartitions()))) - ).andReturn(createdTasks); - replay(activeTaskCreator, standbyTaskCreator); + ).thenReturn(createdTasks); taskManager.handleAssignment( Collections.emptyMap(), mkMap(mkEntry(standbyTaskToBeCreated.id(), standbyTaskToBeCreated.inputPartitions())) ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(tasks).addPendingTaskToInit(createdTasks); + verify(tasks).addPendingTaskToInit(createdTasks); } @Test @@ -551,24 +519,22 @@ public void shouldAssignActiveTaskInTasksRegistryToBeRecycledWithStateUpdaterEna final StandbyTask recycledStandbyTask = standbyTask(taskId03, taskId03ChangelogPartitions) .inState(State.CREATED) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToRecycle)); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - expect(standbyTaskCreator.createStandbyTaskFromActive(activeTaskToRecycle, activeTaskToRecycle.inputPartitions())) - .andReturn(recycledStandbyTask); + when(standbyTaskCreator.createStandbyTaskFromActive(activeTaskToRecycle, activeTaskToRecycle.inputPartitions())) + .thenReturn(recycledStandbyTask); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(activeTaskToRecycle.id()); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( Collections.emptyMap(), mkMap(mkEntry(activeTaskToRecycle.id(), activeTaskToRecycle.inputPartitions())) ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(activeTaskToRecycle).prepareCommit(); - Mockito.verify(tasks).replaceActiveWithStandby(recycledStandbyTask); + verify(activeTaskToRecycle).prepareCommit(); + verify(tasks).replaceActiveWithStandby(recycledStandbyTask); } @Test @@ -576,10 +542,9 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToRecycleIsFoundInTasksRegis final StandbyTask standbyTaskToRecycle = standbyTask(taskId03, taskId03ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToRecycle)); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(activeTaskCreator, standbyTaskCreator); final IllegalStateException illegalStateException = assertThrows( IllegalStateException.class, @@ -590,7 +555,6 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToRecycleIsFoundInTasksRegis ); assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater"); - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -598,20 +562,18 @@ public void shouldAssignActiveTaskInTasksRegistryToBeClosedCleanlyWithStateUpdat final StreamTask activeTaskToClose = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToClose)); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(activeTaskToClose.id()); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment(Collections.emptyMap(), Collections.emptyMap()); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(activeTaskToClose).prepareCommit(); - Mockito.verify(activeTaskToClose).closeClean(); - Mockito.verify(tasks).removeTask(activeTaskToClose); + verify(activeTaskToClose).prepareCommit(); + verify(activeTaskToClose).closeClean(); + verify(tasks).removeTask(activeTaskToClose); } @Test @@ -619,10 +581,9 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToCloseIsFoundInTasksRegistr final StandbyTask standbyTaskToClose = standbyTask(taskId03, taskId03ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToClose)); - replay(activeTaskCreator, standbyTaskCreator); final IllegalStateException illegalStateException = assertThrows( IllegalStateException.class, @@ -630,7 +591,6 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToCloseIsFoundInTasksRegistr ); assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater"); - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -639,21 +599,19 @@ public void shouldAssignActiveTaskInTasksRegistryToUpdateInputPartitionsWithStat .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); final Set newInputPartitions = taskId02Partitions; - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToUpdateInputPartitions)); when(tasks.updateActiveTaskInputPartitions(activeTaskToUpdateInputPartitions, newInputPartitions)).thenReturn(true); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToUpdateInputPartitions.id(), newInputPartitions)), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(activeTaskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(newInputPartitions), any()); + verify(activeTaskToUpdateInputPartitions).updateInputPartitions(eq(newInputPartitions), any()); } @Test @@ -661,19 +619,16 @@ public void shouldResumeActiveRunningTaskInTasksRegistryWithStateUpdaterEnabled( final StreamTask activeTaskToResume = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToResume)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToResume.id(), activeTaskToResume.inputPartitions())), Collections.emptyMap() ); - - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -681,22 +636,20 @@ public void shouldResumeActiveSuspendedTaskInTasksRegistryAndAddToStateUpdater() final StreamTask activeTaskToResume = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.SUSPENDED) .withInputPartitions(taskId03Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToResume)); - expect(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(consumer, Collections.emptyMap())).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToResume.id(), activeTaskToResume.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(activeTaskToResume).resume(); - Mockito.verify(stateUpdater).add(activeTaskToResume); - Mockito.verify(tasks).removeTask(activeTaskToResume); + verify(activeTaskToResume).resume(); + verify(stateUpdater).add(activeTaskToResume); + verify(tasks).removeTask(activeTaskToResume); } @Test @@ -705,10 +658,9 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToUpdateInputPartitionsIsFou .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); final Set newInputPartitions = taskId03Partitions; - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(standbyTaskToUpdateInputPartitions)); - replay(activeTaskCreator, standbyTaskCreator); final IllegalStateException illegalStateException = assertThrows( IllegalStateException.class, @@ -719,7 +671,6 @@ public void shouldThrowDuringAssignmentIfStandbyTaskToUpdateInputPartitionsIsFou ); assertEquals(illegalStateException.getMessage(), "Standby tasks should only be managed by the state updater"); - verify(activeTaskCreator, standbyTaskCreator); } @Test @@ -730,24 +681,22 @@ public void shouldAssignMultipleTasksInTasksRegistryWithStateUpdaterEnabled() { final StreamTask activeTaskToCreate = statefulTask(taskId02, taskId02ChangelogPartitions) .inState(State.CREATED) .withInputPartitions(taskId02Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(activeTaskToClose)); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(activeTaskToClose.id()); - expect(activeTaskCreator.createTasks( + when(activeTaskCreator.createTasks( consumer, mkMap(mkEntry(activeTaskToCreate.id(), activeTaskToCreate.inputPartitions())) - )).andReturn(emptySet()); - expect(standbyTaskCreator.createTasks(Collections.emptyMap())).andReturn(emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + )).thenReturn(emptySet()); + when(standbyTaskCreator.createTasks(Collections.emptyMap())).thenReturn(emptySet()); taskManager.handleAssignment( mkMap(mkEntry(activeTaskToCreate.id(), activeTaskToCreate.inputPartitions())), Collections.emptyMap() ); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(activeTaskToClose).closeClean(); + verify(activeTaskToClose).closeClean(); } @Test @@ -764,10 +713,10 @@ public void shouldAddTasksToStateUpdater() { taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(task00).initializeIfNeeded(); - Mockito.verify(task01).initializeIfNeeded(); - Mockito.verify(stateUpdater).add(task00); - Mockito.verify(stateUpdater).add(task01); + verify(task00).initializeIfNeeded(); + verify(task01).initializeIfNeeded(); + verify(stateUpdater).add(task00); + verify(stateUpdater).add(task01); } @Test @@ -787,10 +736,10 @@ public void shouldRetryInitializationWhenLockExceptionInStateUpdater() { taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(task00).initializeIfNeeded(); - Mockito.verify(task01).initializeIfNeeded(); - Mockito.verify(tasks).addPendingTaskToInit(Collections.singleton(task00)); - Mockito.verify(stateUpdater).add(task01); + verify(task00).initializeIfNeeded(); + verify(task01).initializeIfNeeded(); + verify(tasks).addPendingTaskToInit(Collections.singleton(task00)); + verify(stateUpdater).add(task01); } @Test @@ -811,23 +760,20 @@ public void shouldRecycleTasksRemovedFromStateUpdater() { when(tasks.removePendingTaskToRecycle(task00.id())).thenReturn(taskId00Partitions); when(tasks.removePendingTaskToRecycle(task01.id())).thenReturn(taskId01Partitions); taskManager = setUpTaskManager(StreamsConfigUtils.ProcessingMode.AT_LEAST_ONCE, tasks, true); - expect(activeTaskCreator.createActiveTaskFromStandby(eq(task01), eq(taskId01Partitions), eq(consumer))) - .andStubReturn(task01Converted); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(anyObject()); - expectLastCall().once(); - expect(standbyTaskCreator.createStandbyTaskFromActive(eq(task00), eq(taskId00Partitions))) - .andStubReturn(task00Converted); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createActiveTaskFromStandby(eq(task01), eq(taskId01Partitions), eq(consumer))) + .thenReturn(task01Converted); + when(standbyTaskCreator.createStandbyTaskFromActive(eq(task00), eq(taskId00Partitions))) + .thenReturn(task00Converted); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(task00).suspend(); - Mockito.verify(task01).suspend(); - Mockito.verify(task00Converted).initializeIfNeeded(); - Mockito.verify(task01Converted).initializeIfNeeded(); - Mockito.verify(stateUpdater).add(task00Converted); - Mockito.verify(stateUpdater).add(task01Converted); + verify(task00).suspend(); + verify(task01).suspend(); + verify(task00Converted).initializeIfNeeded(); + verify(task01Converted).initializeIfNeeded(); + verify(stateUpdater).add(task00Converted); + verify(stateUpdater).add(task01Converted); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any()); } @Test @@ -845,17 +791,14 @@ public void shouldCloseTasksRemovedFromStateUpdater() { when(tasks.removePendingTaskToCloseClean(task00.id())).thenReturn(true); when(tasks.removePendingTaskToCloseClean(task01.id())).thenReturn(true); taskManager = setUpTaskManager(StreamsConfigUtils.ProcessingMode.AT_LEAST_ONCE, tasks, true); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(anyObject()); - expectLastCall().once(); - replay(activeTaskCreator); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(activeTaskCreator); - Mockito.verify(task00).suspend(); - Mockito.verify(task00).closeClean(); - Mockito.verify(task01).suspend(); - Mockito.verify(task01).closeClean(); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any()); + verify(task00).suspend(); + verify(task00).closeClean(); + verify(task01).suspend(); + verify(task01).closeClean(); } @Test @@ -873,18 +816,17 @@ public void shouldUpdateInputPartitionsOfTasksRemovedFromStateUpdater() { when(tasks.removePendingTaskToUpdateInputPartitions(task00.id())).thenReturn(taskId02Partitions); when(tasks.removePendingTaskToUpdateInputPartitions(task01.id())).thenReturn(taskId03Partitions); taskManager = setUpTaskManager(StreamsConfigUtils.ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(topologyBuilder); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(task00).updateInputPartitions(Mockito.eq(taskId02Partitions), anyMap()); - Mockito.verify(task00, never()).closeDirty(); - Mockito.verify(task00, never()).closeClean(); - Mockito.verify(stateUpdater).add(task00); - Mockito.verify(task01).updateInputPartitions(Mockito.eq(taskId03Partitions), anyMap()); - Mockito.verify(task01, never()).closeDirty(); - Mockito.verify(task01, never()).closeClean(); - Mockito.verify(stateUpdater).add(task01); + verify(task00).updateInputPartitions(eq(taskId02Partitions), anyMap()); + verify(task00, never()).closeDirty(); + verify(task00, never()).closeClean(); + verify(stateUpdater).add(task00); + verify(task01).updateInputPartitions(eq(taskId03Partitions), anyMap()); + verify(task01, never()).closeDirty(); + verify(task01, never()).closeClean(); + verify(stateUpdater).add(task01); } @Test @@ -899,65 +841,59 @@ public void shouldSuspendRevokedTaskRemovedFromStateUpdater() { when(stateUpdater.hasRemovedTasks()).thenReturn(true); when(stateUpdater.drainRemovedTasks()).thenReturn(mkSet(statefulTask)); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(consumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); - Mockito.verify(statefulTask).suspend(); - Mockito.verify(tasks).addTask(statefulTask); + verify(statefulTask).suspend(); + verify(tasks).addTask(statefulTask); } @Test public void shouldHandleMultipleRemovedTasksFromStateUpdater() { final StreamTask taskToRecycle0 = statefulTask(taskId00, taskId00ChangelogPartitions) - .inState(State.RESTORING) - .withInputPartitions(taskId00Partitions).build(); + .inState(State.RESTORING) + .withInputPartitions(taskId00Partitions).build(); final StandbyTask taskToRecycle1 = standbyTask(taskId01, taskId01ChangelogPartitions) - .inState(State.RUNNING) - .withInputPartitions(taskId01Partitions).build(); + .inState(State.RUNNING) + .withInputPartitions(taskId01Partitions).build(); final StandbyTask convertedTask0 = standbyTask(taskId00, taskId00ChangelogPartitions).build(); final StreamTask convertedTask1 = statefulTask(taskId01, taskId01ChangelogPartitions).build(); final StreamTask taskToClose = statefulTask(taskId02, taskId02ChangelogPartitions) - .inState(State.RESTORING) - .withInputPartitions(taskId02Partitions).build(); + .inState(State.RESTORING) + .withInputPartitions(taskId02Partitions).build(); final StreamTask taskToUpdateInputPartitions = statefulTask(taskId03, taskId03ChangelogPartitions) - .inState(State.RESTORING) - .withInputPartitions(taskId03Partitions).build(); + .inState(State.RESTORING) + .withInputPartitions(taskId03Partitions).build(); when(stateUpdater.hasRemovedTasks()).thenReturn(true); when(stateUpdater.drainRemovedTasks()) - .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions)); + .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions)); when(stateUpdater.restoresActiveTasks()).thenReturn(true); - expect(activeTaskCreator.createActiveTaskFromStandby(eq(taskToRecycle1), eq(taskId01Partitions), eq(consumer))) - .andStubReturn(convertedTask1); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(anyObject()); - expectLastCall().times(2); - expect(standbyTaskCreator.createStandbyTaskFromActive(eq(taskToRecycle0), eq(taskId00Partitions))) - .andStubReturn(convertedTask0); - expect(consumer.assignment()).andReturn(emptySet()).anyTimes(); - consumer.resume(anyObject()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, consumer)) + .thenReturn(convertedTask1); + when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle0, taskId00Partitions)) + .thenReturn(convertedTask0); + when(consumer.assignment()).thenReturn(emptySet()); final TasksRegistry tasks = mock(TasksRegistry.class); when(tasks.removePendingTaskToCloseClean(taskToClose.id())).thenReturn(true); + when(tasks.removePendingTaskToCloseClean(taskToUpdateInputPartitions.id())).thenReturn(false); when(tasks.removePendingTaskToRecycle(taskToRecycle0.id())).thenReturn(taskId00Partitions); when(tasks.removePendingTaskToRecycle(taskToRecycle1.id())).thenReturn(taskId01Partitions); when(tasks.removePendingTaskToRecycle( - argThat(taskId -> !taskId.equals(taskToRecycle0.id()) && !taskId.equals(taskToRecycle1.id()))) + argThat(taskId -> !taskId.equals(taskToRecycle0.id()) && !taskId.equals(taskToRecycle1.id()))) ).thenReturn(null); when(tasks.removePendingTaskToUpdateInputPartitions(taskToUpdateInputPartitions.id())).thenReturn(taskId04Partitions); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); taskManager.setMainConsumer(consumer); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter -> { }); - verify(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer); - Mockito.verify(convertedTask0).initializeIfNeeded(); - Mockito.verify(convertedTask1).initializeIfNeeded(); - Mockito.verify(stateUpdater).add(convertedTask0); - Mockito.verify(stateUpdater).add(convertedTask1); - Mockito.verify(taskToClose).closeClean(); - Mockito.verify(taskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId04Partitions), anyMap()); - Mockito.verify(stateUpdater).add(taskToUpdateInputPartitions); + verify(convertedTask0).initializeIfNeeded(); + verify(convertedTask1).initializeIfNeeded(); + verify(activeTaskCreator, times(2)).closeAndRemoveTaskProducerIfNeeded(any()); + verify(stateUpdater).add(convertedTask0); + verify(stateUpdater).add(convertedTask1); + verify(taskToClose).closeClean(); + verify(taskToUpdateInputPartitions).updateInputPartitions(eq(taskId04Partitions), anyMap()); + verify(stateUpdater).add(taskToUpdateInputPartitions); } @Test @@ -971,8 +907,8 @@ public void shouldAddActiveTaskWithRevokedInputPartitionsInStateUpdaterToPending taskManager.handleRevocation(task.inputPartitions()); - Mockito.verify(tasks).addPendingActiveTaskToSuspend(task.id()); - Mockito.verify(stateUpdater, never()).remove(task.id()); + verify(tasks).addPendingActiveTaskToSuspend(task.id()); + verify(stateUpdater, never()).remove(task.id()); } public void shouldAddMultipleActiveTasksWithRevokedInputPartitionsInStateUpdaterToPendingTasksToSuspend() { @@ -987,8 +923,8 @@ public void shouldAddMultipleActiveTasksWithRevokedInputPartitionsInStateUpdater taskManager.handleRevocation(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - Mockito.verify(tasks).addPendingActiveTaskToSuspend(task1.id()); - Mockito.verify(tasks).addPendingActiveTaskToSuspend(task2.id()); + verify(tasks).addPendingActiveTaskToSuspend(task1.id()); + verify(tasks).addPendingActiveTaskToSuspend(task2.id()); } @Test @@ -1001,8 +937,8 @@ public void shouldNotAddActiveTaskWithoutRevokedInputPartitionsInStateUpdaterToP taskManager.handleRevocation(taskId01Partitions); - Mockito.verify(stateUpdater, never()).remove(task.id()); - Mockito.verify(tasks, never()).addPendingActiveTaskToSuspend(task.id()); + verify(stateUpdater, never()).remove(task.id()); + verify(tasks, never()).addPendingActiveTaskToSuspend(task.id()); } @Test @@ -1015,8 +951,8 @@ public void shouldNotRevokeStandbyTaskInStateUpdaterOnRevocation() { taskManager.handleRevocation(taskId00Partitions); - Mockito.verify(stateUpdater, never()).remove(task.id()); - Mockito.verify(tasks, never()).addPendingActiveTaskToSuspend(task.id()); + verify(stateUpdater, never()).remove(task.id()); + verify(tasks, never()).addPendingActiveTaskToSuspend(task.id()); } @Test @@ -1035,13 +971,13 @@ public void shouldRemoveAllActiveTasksFromStateUpdaterOnPartitionLost() { taskManager.handleLostAll(); - Mockito.verify(stateUpdater).remove(task1.id()); - Mockito.verify(stateUpdater, never()).remove(task2.id()); - Mockito.verify(stateUpdater).remove(task3.id()); - Mockito.verify(tasks).addPendingTaskToCloseDirty(task1.id()); - Mockito.verify(tasks, never()).addPendingTaskToCloseDirty(task2.id()); - Mockito.verify(tasks, never()).addPendingTaskToCloseClean(task2.id()); - Mockito.verify(tasks).addPendingTaskToCloseDirty(task3.id()); + verify(stateUpdater).remove(task1.id()); + verify(stateUpdater, never()).remove(task2.id()); + verify(stateUpdater).remove(task3.id()); + verify(tasks).addPendingTaskToCloseDirty(task1.id()); + verify(tasks, never()).addPendingTaskToCloseDirty(task2.id()); + verify(tasks, never()).addPendingTaskToCloseClean(task2.id()); + verify(tasks).addPendingTaskToCloseDirty(task3.id()); } private TaskManager setupForRevocationAndLost(final Set tasksInStateUpdater, @@ -1060,14 +996,12 @@ public void shouldTransitRestoredTaskToRunning() { final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks); consumer.resume(task.inputPartitions()); - replay(consumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(task).completeRestoration(noOpResetter); - Mockito.verify(task).clearTaskTimeout(); - Mockito.verify(tasks).addTask(task); - verify(consumer); + verify(task).completeRestoration(noOpResetter); + verify(task).clearTaskTimeout(); + verify(tasks).addTask(task); } @Test @@ -1079,14 +1013,12 @@ public void shouldHandleTimeoutExceptionInTransitRestoredTaskToRunning() { final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks); final TimeoutException timeoutException = new TimeoutException(); doThrow(timeoutException).when(task).completeRestoration(noOpResetter); - replay(consumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(task).maybeInitTaskTimeoutOrThrow(anyLong(), Mockito.eq(timeoutException)); - Mockito.verify(tasks, never()).addTask(task); - Mockito.verify(task, never()).clearTaskTimeout(); - verify(consumer); + verify(task).maybeInitTaskTimeoutOrThrow(anyLong(), eq(timeoutException)); + verify(tasks, never()).addTask(task); + verify(task, never()).clearTaskTimeout(); } private TaskManager setUpTransitionToRunningOfRestoredTask(final StreamTask statefulTask, @@ -1119,17 +1051,15 @@ public void shouldRecycleRestoredTask() { .inState(State.CREATED) .withInputPartitions(taskId00Partitions).build(); final TaskManager taskManager = setUpRecycleRestoredTask(statefulTask); - expect(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) - .andStubReturn(standbyTask); + when(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) + .thenReturn(standbyTask); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask.id()); - replay(activeTaskCreator, standbyTaskCreator); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(activeTaskCreator, standbyTaskCreator); - Mockito.verify(statefulTask).suspend(); - Mockito.verify(standbyTask).initializeIfNeeded(); - Mockito.verify(stateUpdater).add(standbyTask); + verify(statefulTask).suspend(); + verify(standbyTask).initializeIfNeeded(); + verify(stateUpdater).add(standbyTask); } @Test @@ -1138,18 +1068,16 @@ public void shouldHandleExceptionThrownDuringConversionInRecycleRestoredTask() { .inState(State.RESTORING) .withInputPartitions(taskId00Partitions).build(); final TaskManager taskManager = setUpRecycleRestoredTask(statefulTask); - expect(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) - .andThrow(new RuntimeException()); - replay(standbyTaskCreator); + when(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) + .thenThrow(new RuntimeException()); assertThrows( StreamsException.class, () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter) ); - verify(standbyTaskCreator); - Mockito.verify(stateUpdater, never()).add(any()); - Mockito.verify(statefulTask).closeDirty(); + verify(stateUpdater, never()).add(any()); + verify(statefulTask).closeDirty(); } @Test @@ -1161,19 +1089,17 @@ public void shouldHandleExceptionThrownDuringTaskInitInRecycleRestoredTask() { .inState(State.CREATED) .withInputPartitions(taskId00Partitions).build(); final TaskManager taskManager = setUpRecycleRestoredTask(statefulTask); - expect(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) - .andStubReturn(standbyTask); + when(standbyTaskCreator.createStandbyTaskFromActive(statefulTask, statefulTask.inputPartitions())) + .thenReturn(standbyTask); doThrow(StreamsException.class).when(standbyTask).initializeIfNeeded(); - replay(standbyTaskCreator); assertThrows( StreamsException.class, () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter) ); - verify(standbyTaskCreator); - Mockito.verify(stateUpdater, never()).add(any()); - Mockito.verify(standbyTask).closeDirty(); + verify(stateUpdater, never()).add(any()); + verify(standbyTask).closeDirty(); } private TaskManager setUpRecycleRestoredTask(final StreamTask statefulTask) { @@ -1193,15 +1119,13 @@ public void shouldCloseCleanRestoredTask() { final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpCloseCleanRestoredTask(statefulTask, tasks); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask.id()); - replay(activeTaskCreator); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(activeTaskCreator); - Mockito.verify(statefulTask).suspend(); - Mockito.verify(statefulTask).closeClean(); - Mockito.verify(statefulTask, never()).closeDirty(); - Mockito.verify(tasks, never()).removeTask(statefulTask); + verify(statefulTask).suspend(); + verify(statefulTask).closeClean(); + verify(statefulTask, never()).closeDirty(); + verify(tasks, never()).removeTask(statefulTask); } @Test @@ -1213,16 +1137,14 @@ public void shouldHandleExceptionThrownDuringCloseInCloseCleanRestoredTask() { final TaskManager taskManager = setUpCloseCleanRestoredTask(statefulTask, tasks); doThrow(RuntimeException.class).when(statefulTask).closeClean(); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask.id()); - replay(activeTaskCreator); assertThrows( RuntimeException.class, () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter) ); - verify(activeTaskCreator); - Mockito.verify(statefulTask).closeDirty(); - Mockito.verify(tasks, never()).removeTask(statefulTask); + verify(statefulTask).closeDirty(); + verify(tasks, never()).removeTask(statefulTask); } @Test @@ -1230,20 +1152,18 @@ public void shouldHandleExceptionThrownDuringClosingTaskProducerInCloseCleanRest final StreamTask statefulTask = statefulTask(taskId00, taskId00ChangelogPartitions) .inState(State.CLOSED) .withInputPartitions(taskId00Partitions).build(); + final TaskId taskId = statefulTask.id(); final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpCloseCleanRestoredTask(statefulTask, tasks); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask.id()); - expectLastCall().andThrow(new RuntimeException("Something happened")); - replay(activeTaskCreator); + doThrow(new RuntimeException("Something happened")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId); assertThrows( RuntimeException.class, () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter) ); - verify(activeTaskCreator); - Mockito.verify(statefulTask, never()).closeDirty(); - Mockito.verify(tasks, never()).removeTask(statefulTask); + verify(statefulTask, never()).closeDirty(); + verify(tasks, never()).removeTask(statefulTask); } private TaskManager setUpCloseCleanRestoredTask(final StreamTask statefulTask, @@ -1268,16 +1188,14 @@ public void shouldCloseDirtyRestoredTask() { when(stateUpdater.restoresActiveTasks()).thenReturn(true); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask.id()); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(activeTaskCreator); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(activeTaskCreator); - Mockito.verify(statefulTask).prepareCommit(); - Mockito.verify(statefulTask).suspend(); - Mockito.verify(statefulTask).closeDirty(); - Mockito.verify(statefulTask, never()).closeClean(); - Mockito.verify(tasks, never()).removeTask(statefulTask); + verify(statefulTask).prepareCommit(); + verify(statefulTask).suspend(); + verify(statefulTask).closeDirty(); + verify(statefulTask, never()).closeClean(); + verify(tasks, never()).removeTask(statefulTask); } @Test @@ -1292,15 +1210,13 @@ public void shouldUpdateInputPartitionsOfRestoredTask() { when(stateUpdater.restoresActiveTasks()).thenReturn(true); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); consumer.resume(statefulTask.inputPartitions()); - replay(consumer, topologyBuilder); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); - Mockito.verify(statefulTask).updateInputPartitions(Mockito.eq(taskId01Partitions), anyMap()); - Mockito.verify(statefulTask).completeRestoration(noOpResetter); - Mockito.verify(statefulTask).clearTaskTimeout(); - Mockito.verify(tasks).addTask(statefulTask); + verify(statefulTask).updateInputPartitions(eq(taskId01Partitions), anyMap()); + verify(statefulTask).completeRestoration(noOpResetter); + verify(statefulTask).clearTaskTimeout(); + verify(tasks).addTask(statefulTask); } @Test @@ -1315,13 +1231,11 @@ public void shouldSuspendRestoredTaskIfRevoked() { when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(mkSet(statefulTask)); when(stateUpdater.restoresActiveTasks()).thenReturn(true); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(consumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); - Mockito.verify(statefulTask).suspend(); - Mockito.verify(tasks).addTask(statefulTask); + verify(statefulTask).suspend(); + verify(tasks).addTask(statefulTask); } @Test @@ -1346,8 +1260,8 @@ public void shouldHandleMultipleRestoredTasks() { .withInputPartitions(taskId04Partitions).build(); final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - expect(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle, taskToRecycle.inputPartitions())) - .andStubReturn(recycledStandbyTask); + when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle, taskToRecycle.inputPartitions())) + .thenReturn(recycledStandbyTask); when(tasks.removePendingTaskToRecycle(taskToRecycle.id())).thenReturn(taskId01Partitions); when(tasks.removePendingTaskToRecycle( argThat(taskId -> !taskId.equals(taskToRecycle.id()))) @@ -1372,16 +1286,15 @@ public void shouldHandleMultipleRestoredTasks() { taskToCloseDirty, taskToUpdateInputPartitions )); - replay(standbyTaskCreator, topologyBuilder); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - Mockito.verify(tasks).addTask(taskToTransitToRunning); - Mockito.verify(stateUpdater).add(recycledStandbyTask); - Mockito.verify(stateUpdater).add(recycledStandbyTask); - Mockito.verify(taskToCloseClean).closeClean(); - Mockito.verify(taskToCloseDirty).closeDirty(); - Mockito.verify(taskToUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId05Partitions), anyMap()); + verify(tasks).addTask(taskToTransitToRunning); + verify(stateUpdater).add(recycledStandbyTask); + verify(stateUpdater).add(recycledStandbyTask); + verify(taskToCloseClean).closeClean(); + verify(taskToCloseDirty).closeDirty(); + verify(taskToUpdateInputPartitions).updateInputPartitions(eq(taskId05Partitions), anyMap()); } @Test @@ -1488,9 +1401,9 @@ public void shouldRethrowTaskCorruptedExceptionFromInitialization() { () -> taskManager.checkStateUpdater(time.milliseconds(), noOpResetter) ); - Mockito.verify(tasks).addTask(statefulTask0); - Mockito.verify(tasks).addTask(statefulTask1); - Mockito.verify(stateUpdater).add(statefulTask2); + verify(tasks).addTask(statefulTask0); + verify(tasks).addTask(statefulTask1); + verify(stateUpdater).add(statefulTask2); assertEquals(mkSet(taskId00, taskId01), thrown.corruptedTasks()); assertEquals("Tasks [0_1, 0_0] are corrupted and hence need to be re-initialized", thrown.getMessage()); } @@ -1500,27 +1413,20 @@ public void shouldIdempotentlyUpdateSubscriptionFromActiveAssignment() { final TopicPartition newTopicPartition = new TopicPartition("topic2", 1); final Map> assignment = mkMap(mkEntry(taskId01, mkSet(t1p1, newTopicPartition))); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(emptyList()); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - - topologyBuilder.addSubscribedTopicsFromAssignment(eq(asList(t1p1, newTopicPartition)), anyString()); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(emptyList()); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(assignment, emptyMap()); - verify(activeTaskCreator, topologyBuilder); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(eq(asList(t1p1, newTopicPartition)), anyString()); } @Test public void shouldNotLockAnythingIfStateDirIsEmpty() { - expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()).once(); + when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>()); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); - verify(stateDirectory); assertTrue(taskManager.lockedTaskDirectories().isEmpty()); } @@ -1534,23 +1440,19 @@ public void shouldTryToLockValidTaskDirsAtRebalanceStart() throws Exception { taskId10.toString(), "dummy" ); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); - verify(stateDirectory); assertThat(taskManager.lockedTaskDirectories(), is(singleton(taskId01))); } @Test public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() { final Set assigned = mkSet(t1p0, t1p1); - expect(consumer.assignment()).andReturn(assigned); - consumer.pause(assigned); - replay(consumer); + when(consumer.assignment()).thenReturn(assigned); taskManager.handleRebalanceComplete(); - verify(consumer); + verify(consumer).pause(assigned); } @Test @@ -1558,42 +1460,41 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() { final StreamTask statefulTask0 = statefulTask(taskId00, taskId00ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId00Partitions).build(); - final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); when(tasks.allTasks()).thenReturn(mkSet(statefulTask0)); final Set assigned = mkSet(t1p0, t1p1); - expect(consumer.assignment()).andReturn(assigned); - consumer.pause(mkSet(t1p1)); - replay(consumer); + when(consumer.assignment()).thenReturn(assigned); taskManager.handleRebalanceComplete(); - verify(consumer); + verify(consumer).pause(mkSet(t1p1)); } @Test public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception { expectLockObtainedFor(taskId00, taskId01, taskId02); - expectUnlockFor(taskId02); makeTaskFolders( taskId00.toString(), // active task taskId01.toString(), // standby task taskId02.toString() // unassigned but able to lock ); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02))); handleAssignment(taskId00Assignment, taskId01Assignment, emptyMap()); - reset(consumer); - expectConsumerAssignmentPaused(consumer); - replay(consumer); + + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(consumer.assignment()).thenReturn(assignment); taskManager.handleRebalanceComplete(); assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01))); - verify(stateDirectory); + verify(consumer).pause(assignment); + verify(stateDirectory).unlock(taskId02); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -1633,7 +1534,6 @@ private void computeOffsetSumAndVerify(final Map changelog final Map expectedOffsetSums) throws Exception { expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); final StateMachineTask restoringTask = handleAssignment( @@ -1656,7 +1556,6 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception { expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); final StateMachineTask restoringTask = handleAssignment( @@ -1667,6 +1566,8 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception { restoringTask.setChangelogOffsets(changelogOffsets); assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -1681,7 +1582,6 @@ public void shouldComputeOffsetSumForUnassignedTaskWeCanLock() throws Exception makeTaskFolders(taskId00.toString()); writeCheckpointFile(taskId00, changelogOffsets); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums)); @@ -1698,13 +1598,11 @@ public void shouldComputeOffsetSumFromCheckpointFileForUninitializedTask() throw expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); writeCheckpointFile(taskId00, changelogOffsets); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); final StateMachineTask uninitializedTask = new StateMachineTask(taskId00, taskId00Partitions, true); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singleton(uninitializedTask)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singleton(uninitializedTask)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(uninitializedTask.state(), is(State.CREATED)); @@ -1723,14 +1621,12 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); writeCheckpointFile(taskId00, changelogOffsets); - replay(stateDirectory); final StateMachineTask closedTask = new StateMachineTask(taskId00, taskId00Partitions, true); taskManager.handleRebalanceStart(singleton("topic")); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singleton(closedTask)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singleton(closedTask)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); closedTask.suspend(); @@ -1744,7 +1640,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep public void shouldNotReportOffsetSumsForTaskWeCantLock() throws Exception { expectLockFailedFor(taskId00); makeTaskFolders(taskId00.toString()); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); assertTrue(taskManager.lockedTaskDirectories().isEmpty()); @@ -1755,8 +1650,7 @@ public void shouldNotReportOffsetSumsForTaskWeCantLock() throws Exception { public void shouldNotReportOffsetSumsAndReleaseLockForUnassignedTaskWithoutCheckpoint() throws Exception { expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); - expect(stateDirectory.checkpointFileFor(taskId00)).andReturn(getCheckpointFile(taskId00)); - replay(stateDirectory); + when(stateDirectory.checkpointFileFor(taskId00)).thenReturn(getCheckpointFile(taskId00)); taskManager.handleRebalanceStart(singleton("topic")); assertTrue(taskManager.getTaskOffsetSums().isEmpty()); @@ -1776,7 +1670,6 @@ public void shouldPinOffsetSumToLongMaxValueInCaseOfOverflow() throws Exception expectLockObtainedFor(taskId00); makeTaskFolders(taskId00.toString()); writeCheckpointFile(taskId00, changelogOffsets); - replay(stateDirectory); taskManager.handleRebalanceStart(singleton("topic")); assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums)); @@ -1790,23 +1683,9 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() { // first `handleAssignment` expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(emptyMap()))).andStubReturn(emptyList()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(emptyList()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - - // `handleRevocation` - consumer.commitSync(offsets); - expectLastCall(); - - // second `handleAssignment` - consumer.commitSync(offsets); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), eq(emptyMap()))).thenReturn(emptyList()); + when(standbyTaskCreator.createTasks(any())).thenReturn(emptyList()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -1832,30 +1711,26 @@ public void closeClean() { // first `handleAssignment` expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(emptyMap()))).andStubReturn(emptyList()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(emptyList()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), eq(emptyMap()))).thenReturn(emptyList()); + when(standbyTaskCreator.createTasks(any())).thenReturn(emptyList()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); taskManager.handleRevocation(taskId00Partitions); final RuntimeException thrown = assertThrows( - RuntimeException.class, - () -> taskManager.handleAssignment(emptyMap(), emptyMap()) + RuntimeException.class, + () -> taskManager.handleAssignment(emptyMap(), emptyMap()) ); assertThat(task00.state(), is(Task.State.CLOSED)); assertThat( - thrown.getMessage(), - is("Encounter unexpected fatal error for task 0_0") + thrown.getMessage(), + is("Encounter unexpected fatal error for task 0_0") ); assertThat(thrown.getCause().getMessage(), is("KABOOM!")); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); + verify(topologyBuilder, times(2)).addSubscribedTopicsFromAssignment(any(), anyString()); } @Test @@ -1865,28 +1740,16 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { // `handleAssignment` expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))).andStubReturn(singletonList(task01)); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))).thenReturn(singletonList(task01)); + // The second attempt will return empty tasks. makeTaskFolders(taskId00.toString(), taskId01.toString()); expectLockObtainedFor(taskId00, taskId01); - // The second attempt will return empty tasks. - makeTaskFolders(); - expectLockObtainedFor(); - replay(stateDirectory); - taskManager.handleRebalanceStart(emptySet()); assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01))); - // `handleLostAll` - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); - taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); @@ -1905,20 +1768,21 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { taskManager.handleRebalanceStart(emptySet()); assertThat(taskManager.lockedTaskDirectories(), is(emptySet())); + + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + // `handleLostAll` + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); + + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldReInitializeThreadProducerOnHandleLostAllIfEosV2Enabled() { - activeTaskCreator.reInitializeThreadProducer(); - expectLastCall(); - final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false); - replay(activeTaskCreator); - taskManager.handleLostAll(); - verify(activeTaskCreator); + verify(activeTaskCreator).reInitializeThreadProducer(); } @Test @@ -1929,18 +1793,10 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { // `handleAssignment` expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(emptyList()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(emptyList()); - // `handleAssignment` - consumer.commitSync(offsets); - expectLastCall(); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall().andThrow(new RuntimeException("KABOOM!")); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + doThrow(new RuntimeException("KABOOM!")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -1959,14 +1815,14 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { ); assertThat(thrown.getCause(), instanceOf(RuntimeException.class)); assertThat(thrown.getCause().getMessage(), is("KABOOM!")); + verify(topologyBuilder, times(2)).addSubscribedTopicsFromAssignment(any(), anyString()); + + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldReviveCorruptTasks() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - EasyMock.expectLastCall().once(); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final AtomicBoolean enforcedCheckpoint = new AtomicBoolean(false); final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) { @@ -1981,12 +1837,9 @@ public void postCommit(final boolean enforceCheckpoint) { // `handleAssignment` expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); + when(consumer.assignment()).thenReturn(taskId00Partitions); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2002,15 +1855,15 @@ public void postCommit(final boolean enforceCheckpoint) { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(stateManager); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verify(consumer).resume(taskId00Partitions); } @Test public void shouldReviveCorruptTasksEvenIfTheyCannotCloseClean() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) { @Override @@ -2021,12 +1874,9 @@ public void suspend() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); + when(consumer.assignment()).thenReturn(taskId00Partitions); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2040,15 +1890,15 @@ public void suspend() { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(stateManager); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verify(consumer).resume(taskId00Partitions); } @Test public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final StateMachineTask nonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); @@ -2057,18 +1907,14 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { assignment.putAll(taskId01Assignment); // `handleAssignment` - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))) - .andStubReturn(asList(corruptedTask, nonCorruptedTask)); - expect(standbyTaskCreator.createTasks(anyObject())) - .andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(assignment))) + .thenReturn(asList(corruptedTask, nonCorruptedTask)); + when(standbyTaskCreator.createTasks(any())) + .thenReturn(Collections.emptySet()); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(consumer.assignment()).andReturn(taskId00Partitions); + when(consumer.assignment()).thenReturn(taskId00Partitions); // check that we should not commit empty map either - consumer.commitSync(eq(emptyMap())); - expectLastCall().andStubThrow(new AssertionError("should not invoke commitSync when offset map is empty")); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + doThrow(new AssertionError("should not invoke commitSync when offset map is empty")).when(consumer).commitSync(eq(emptyMap())); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2083,14 +1929,15 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { assertThat(nonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet())); assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions)); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verify(consumer).resume(taskId00Partitions); } @Test public void shouldNotCommitNonRunningNonCorruptedTasks() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final StateMachineTask nonRunningNonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); @@ -2101,14 +1948,12 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() { assignment.putAll(taskId01Assignment); // `handleAssignment` - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))) - .andStubReturn(asList(corruptedTask, nonRunningNonCorruptedTask)); - expect(standbyTaskCreator.createTasks(anyObject())) - .andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignment))) + .thenReturn(asList(corruptedTask, nonRunningNonCorruptedTask)); + when(standbyTaskCreator.createTasks(any())) + .thenReturn(Collections.emptySet()); + + when(consumer.assignment()).thenReturn(taskId00Partitions); taskManager.handleAssignment(assignment, emptyMap()); @@ -2119,16 +1964,14 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() { assertThat(nonRunningNonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet())); assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions)); - verify(activeTaskCreator); assertFalse(nonRunningNonCorruptedTask.commitPrepared); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); } @Test public void shouldCleanAndReviveCorruptedStandbyTasksBeforeCommittingNonCorruptedTasks() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask corruptedStandby = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); final StateMachineTask runningNonCorruptedActive = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) { @@ -2139,15 +1982,11 @@ public Map prepareCommit() { }; // handleAssignment - expect(standbyTaskCreator.createTasks(eq(taskId00Assignment))).andStubReturn(singleton(corruptedStandby)); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId01Assignment))).andStubReturn(singleton(runningNonCorruptedActive)); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(standbyTaskCreator.createTasks(eq(taskId00Assignment))).thenReturn(singleton(corruptedStandby)); + when(activeTaskCreator.createTasks(any(), eq(taskId01Assignment))).thenReturn(singleton(runningNonCorruptedActive)); expectRestoreToBeCompleted(consumer, changeLogReader); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); - taskManager.handleAssignment(taskId01Assignment, taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2163,13 +2002,16 @@ public Map prepareCommit() { assertThat(corruptedStandby.commitPrepared, is(true)); assertThat(corruptedStandby.state(), is(Task.State.CREATED)); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { - final ProcessorStateManager stateManager = EasyMock.createNiceMock(ProcessorStateManager.class); - expect(stateDirectory.listNonEmptyTaskDirectories()).andStubReturn(new ArrayList<>()); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); + when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>()); final StateMachineTask corruptedActive = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); @@ -2182,18 +2024,12 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { final Map> assignment = new HashMap<>(); assignment.putAll(taskId00Assignment); assignment.putAll(taskId01Assignment); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(asList(corruptedActive, uncorruptedActive)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); - topologyBuilder.addSubscribedTopicsFromMetadata(eq(singleton(topic1)), anyObject()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(corruptedActive, uncorruptedActive)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader, stateDirectory, stateManager); + when(consumer.assignment()).thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); uncorruptedActive.setCommittableOffsetsAndMetadata(offsets); @@ -2215,14 +2051,15 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { assertThat(uncorruptedActive.commitCompleted, is(false)); assertThat(uncorruptedActive.state(), is(State.RUNNING)); - verify(consumer); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + verify(topologyBuilder).addSubscribedTopicsFromMetadata(eq(singleton(topic1)), any()); + + verify(consumer).resume(union(HashSet::new, taskId00Partitions, taskId01Partitions)); } @Test public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCommitWithALOS() { - final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class); - stateManager.markChangelogAsCorrupted(taskId00Partitions); - replay(stateManager); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask corruptedActive = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final StateMachineTask uncorruptedActive = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) { @@ -2238,19 +2075,14 @@ public void markChangelogAsCorrupted(final Collection partitions final Map> assignment = new HashMap<>(); assignment.putAll(taskId00Assignment); assignment.putAll(taskId01Assignment); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(asList(corruptedActive, uncorruptedActive)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(corruptedActive, uncorruptedActive)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); expectRestoreToBeCompleted(consumer, changeLogReader); - consumer.commitSync(offsets); - expectLastCall().andThrow(new TimeoutException()); - - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); + doThrow(new TimeoutException()).when(consumer).commitSync(offsets); - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader); + when(consumer.assignment()).thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2280,15 +2112,18 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(corruptedActive.state(), is(Task.State.CREATED)); assertThat(uncorruptedActive.state(), is(Task.State.CREATED)); - verify(consumer); + verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verify(consumer).resume(union(HashSet::new, taskId00Partitions, taskId01Partitions)); } @Test public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCommitDuringHandleCorruptedWithEOS() { final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false); final StreamsProducer producer = mock(StreamsProducer.class); - expect(activeTaskCreator.threadProducer()).andStubReturn(producer); - final ProcessorStateManager stateManager = EasyMock.createMock(ProcessorStateManager.class); + when(activeTaskCreator.threadProducer()).thenReturn(producer); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final AtomicBoolean corruptedTaskChangelogMarkedAsCorrupted = new AtomicBoolean(false); final StateMachineTask corruptedActiveTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) { @@ -2316,21 +2151,17 @@ public void markChangelogAsCorrupted(final Collection partitions final Map> assignment = new HashMap<>(); assignment.putAll(taskId00Assignment); assignment.putAll(taskId01Assignment); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(asList(corruptedActiveTask, uncorruptedActiveTask)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString()); - expectLastCall().anyTimes(); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); expectRestoreToBeCompleted(consumer, changeLogReader); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); + when(consumer.groupMetadata()).thenReturn(groupMetadata); doThrow(new TimeoutException()).when(producer).commitTransaction(offsets, groupMetadata); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - - replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader, stateManager); + when(consumer.assignment()).thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2366,7 +2197,9 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(uncorruptedActiveTask.state(), is(Task.State.CREATED)); assertThat(corruptedTaskChangelogMarkedAsCorrupted.get(), is(true)); assertThat(uncorruptedTaskChangelogMarkedAsCorrupted.get(), is(true)); - verify(consumer); + verify(topologyBuilder).addSubscribedTopicsFromAssignment(any(), anyString()); + + verify(consumer).resume(union(HashSet::new, taskId00Partitions, taskId01Partitions)); } @Test @@ -2393,22 +2226,17 @@ public void markChangelogAsCorrupted(final Collection partitions expectedCommittedOffsets.putAll(offsets01); final Map> assignmentActive = mkMap( - mkEntry(taskId00, taskId00Partitions), - mkEntry(taskId01, taskId01Partitions), - mkEntry(taskId02, taskId02Partitions) + mkEntry(taskId00, taskId00Partitions), + mkEntry(taskId01, taskId01Partitions), + mkEntry(taskId02, taskId02Partitions) ); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))).andReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall().andThrow(new TimeoutException()); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); + doThrow(new TimeoutException()).when(consumer).commitSync(expectedCommittedOffsets); + when(consumer.assignment()).thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); taskManager.handleAssignment(assignmentActive, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2421,14 +2249,16 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(revokedActiveTask.state(), is(State.SUSPENDED)); assertThat(unrevokedActiveTaskWithCommitNeeded.state(), is(State.CREATED)); assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); + + verify(consumer).resume(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); } @Test public void shouldCloseAndReviveUncorruptedTasksWhenTimeoutExceptionThrownFromCommitDuringRevocationWithEOS() { final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false); final StreamsProducer producer = mock(StreamsProducer.class); - expect(activeTaskCreator.threadProducer()).andStubReturn(producer); - final ProcessorStateManager stateManager = EasyMock.createMock(ProcessorStateManager.class); + when(activeTaskCreator.threadProducer()).thenReturn(producer); + final ProcessorStateManager stateManager = mock(ProcessorStateManager.class); final StateMachineTask revokedActiveTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final Map revokedActiveTaskOffsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); @@ -2464,19 +2294,15 @@ public void markChangelogAsCorrupted(final Collection partitions expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))).andReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); + when(consumer.groupMetadata()).thenReturn(groupMetadata); doThrow(new TimeoutException()).when(producer).commitTransaction(expectedCommittedOffsets, groupMetadata); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader, stateManager); + when(consumer.assignment()).thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); taskManager.handleAssignment(assignmentActive, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2495,6 +2321,8 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(revokedActiveTask.state(), is(State.SUSPENDED)); assertThat(unrevokedActiveTask.state(), is(State.CREATED)); assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); + + verify(consumer).resume(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); } @Test @@ -2502,12 +2330,9 @@ public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(standbyTaskCreator.createTasks(eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(activeTaskCreator.createTasks(anyObject(), anyObject())).andStubReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).andStubReturn(Collections.emptySet()); - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(standbyTaskCreator.createTasks(eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(emptyMap(), taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2517,6 +2342,8 @@ public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { assertThat(task00.state(), is(Task.State.CLOSED)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2527,13 +2354,10 @@ public void shouldAddNonResumedSuspendedTasks() { expectRestoreToBeCompleted(consumer, changeLogReader); // expect these calls twice (because we're going to tryToCompleteRestoration twice) expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))).andReturn(singletonList(task01)).anyTimes(); - expect(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(eq(asList(t1p0)), anyString()); - expectLastCall().anyTimes(); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader, topologyBuilder); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))).thenReturn(singletonList(task01)); + when(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2545,7 +2369,9 @@ public void shouldAddNonResumedSuspendedTasks() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThat(task01.state(), is(Task.State.RUNNING)); - verify(activeTaskCreator); + verify(topologyBuilder, times(2)).addSubscribedTopicsFromAssignment(eq(asList(t1p0)), anyString()); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2553,12 +2379,9 @@ public void shouldUpdateInputPartitionsAfterRebalance() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true); expectRestoreToBeCompleted(consumer, changeLogReader); - // expect these calls twice (because we're going to tryToCompleteRestoration twice) - expectRestoreToBeCompleted(consumer, changeLogReader, false); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2570,7 +2393,8 @@ public void shouldUpdateInputPartitionsAfterRebalance() { assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); assertEquals(newPartitionsSet, task00.inputPartitions()); - verify(activeTaskCreator, consumer, changeLogReader); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2578,15 +2402,10 @@ public void shouldAddNewActiveTasks() { final Map> assignment = taskId00Assignment; final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); - changeLogReader.enforceRestoreActive(); - expectLastCall(); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(consumer, activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(consumer.assignment()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -2597,14 +2416,14 @@ public void shouldAddNewActiveTasks() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(activeTaskCreator); + verify(consumer).resume(eq(emptySet())); } @Test public void shouldNotCompleteRestorationIfTasksCannotInitialize() { final Map> assignment = mkMap( - mkEntry(taskId00, taskId00Partitions), - mkEntry(taskId01, taskId01Partitions) + mkEntry(taskId00, taskId00Partitions), + mkEntry(taskId01, taskId01Partitions) ); final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true) { @Override @@ -2619,17 +2438,10 @@ public void initializeIfNeeded() { } }; - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); - changeLogReader.enforceRestoreActive(); - expectLastCall(); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(asList(task00, task01)); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(consumer, activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(consumer.assignment()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01)); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -2641,17 +2453,16 @@ public void initializeIfNeeded() { assertThat(task00.state(), is(Task.State.CREATED)); assertThat(task01.state(), is(Task.State.CREATED)); assertThat( - taskManager.activeTaskMap(), - Matchers.equalTo(mkMap(mkEntry(taskId00, task00), mkEntry(taskId01, task01))) + taskManager.activeTaskMap(), + Matchers.equalTo(mkMap(mkEntry(taskId00, task00), mkEntry(taskId01, task01))) ); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(activeTaskCreator); } @Test public void shouldNotCompleteRestorationIfTaskCannotCompleteRestoration() { final Map> assignment = mkMap( - mkEntry(taskId00, taskId00Partitions) + mkEntry(taskId00, taskId00Partitions) ); final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true) { @Override @@ -2660,17 +2471,10 @@ public void completeRestoration(final java.util.function.Consumer> assignmentActive = mkMap( - mkEntry(taskId00, taskId00Partitions), - mkEntry(taskId01, taskId01Partitions), - mkEntry(taskId02, taskId02Partitions) + mkEntry(taskId00, taskId00Partitions), + mkEntry(taskId01, taskId01Partitions), + mkEntry(taskId02, taskId02Partitions) ); final Map> assignmentStandby = mkMap( - mkEntry(taskId10, taskId10Partitions) + mkEntry(taskId10, taskId10Partitions) ); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))) - .andReturn(asList(task00, task01, task02)); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))) + .thenReturn(asList(task00, task01, task02)); - expect(activeTaskCreator.threadProducer()).andReturn(producer); + when(activeTaskCreator.threadProducer()).thenReturn(producer); activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))) - .andReturn(singletonList(task10)); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))) + .thenReturn(singletonList(task10)); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); - producer.commitTransaction(expectedCommittedOffsets, groupMetadata); - expectLastCall(); - - task00.committedOffsets(); - EasyMock.expectLastCall(); - task01.committedOffsets(); - EasyMock.expectLastCall(); - task02.committedOffsets(); - EasyMock.expectLastCall(); - task10.committedOffsets(); - EasyMock.expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(consumer.groupMetadata()).thenReturn(groupMetadata); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2782,6 +2570,9 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo assertThat(task01.commitNeeded, is(false)); assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); + verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2817,16 +2608,10 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { ); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))) - .andReturn(asList(task00, task01, task02)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall(); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))) - .andReturn(singletonList(task10)); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))) + .thenReturn(asList(task00, task01, task02)); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))) + .thenReturn(singletonList(task10)); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2843,6 +2628,9 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { assertThat(task01.commitPrepared, is(true)); assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); + verify(consumer).commitSync(expectedCommittedOffsets); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2859,14 +2647,10 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))).andReturn(singleton(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))).andReturn(singletonList(task10)); - expect(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - topologyBuilder.addSubscribedTopicsFromAssignment(eq(asList(t1p0)), anyString()); - expectLastCall().anyTimes(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader, topologyBuilder); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(singleton(task00)); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))).thenReturn(singletonList(task10)); + when(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2877,6 +2661,9 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { assertThat(task00.commitNeeded, is(true)); assertThat(task10.commitPrepared, is(false)); + verify(topologyBuilder, times(2)).addSubscribedTopicsFromAssignment(eq(asList(t1p0)), anyString()); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -2893,12 +2680,10 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))).andReturn(singleton(task00)); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))).andReturn(singletonList(task10)); - expect(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))).thenReturn(singleton(task00)); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))).thenReturn(singletonList(task10)); + when(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2914,12 +2699,9 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { public void shouldNotCommitCreatedTasksOnRevocationOrClosure() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); - expectLastCall().once(); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(task00.state(), is(Task.State.CREATED)); @@ -2929,6 +2711,7 @@ public void shouldNotCommitCreatedTasksOnRevocationOrClosure() { taskManager.handleAssignment(emptyMap(), emptyMap()); assertThat(task00.state(), is(Task.State.CLOSED)); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); } @Test @@ -2942,9 +2725,8 @@ public void suspend() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); @@ -2952,7 +2734,7 @@ public void suspend() { assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions)); assertThat(task00.state(), is(Task.State.SUSPENDED)); - verify(consumer); + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3013,17 +2795,10 @@ public void closeDirty() { } }; - resetToStrict(changeLogReader); - changeLogReader.enforceRestoreActive(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))) - .andStubReturn(asList(task00, task01, task02, task03)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(anyObject()); - expectLastCall().times(4); - activeTaskCreator.closeThreadProducerIfNeeded(); - expectLastCall(); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))) + .thenReturn(asList(task00, task01, task02, task03)); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -3067,7 +2842,8 @@ public void closeDirty() { assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); // the active task creator should also get closed (so that it closes the thread producer if applicable) - verify(activeTaskCreator, changeLogReader); + verify(activeTaskCreator, times(4)).closeAndRemoveTaskProducerIfNeeded(any()); + verify(activeTaskCreator).closeThreadProducerIfNeeded(); } @Test @@ -3085,16 +2861,10 @@ public Set changelogPartitions() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task00.setCommittableOffsetsAndMetadata(offsets); - resetToStrict(changeLogReader); - changeLogReader.enforceRestoreActive(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(singletonList(task00)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); - expectLastCall().andThrow(new RuntimeException("whatever")); - activeTaskCreator.closeThreadProducerIfNeeded(); - expectLastCall(); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00)); + doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -3120,7 +2890,7 @@ public Set changelogPartitions() { assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); // the active task creator should also get closed (so that it closes the thread producer if applicable) - verify(activeTaskCreator, changeLogReader); + verify(activeTaskCreator).closeThreadProducerIfNeeded(); } @Test @@ -3136,16 +2906,10 @@ public Set changelogPartitions() { } }; - resetToStrict(changeLogReader); - changeLogReader.enforceRestoreActive(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(singletonList(task00)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); - expectLastCall(); - activeTaskCreator.closeThreadProducerIfNeeded(); - expectLastCall().andThrow(new RuntimeException("whatever")); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00)); + doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeThreadProducerIfNeeded(); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -3171,7 +2935,7 @@ public Set changelogPartitions() { assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); // the active task creator should also get closed (so that it closes the thread producer if applicable) - verify(activeTaskCreator, changeLogReader); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(eq(taskId00)); } @Test @@ -3223,8 +2987,6 @@ public void suspend() { taskManager.addTask(task01); taskManager.addTask(task02); - replay(activeTaskCreator); - final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(union(HashSet::new, taskId01Partitions, taskId02Partitions))); assertThat(thrown.getCause().getMessage(), is("task 0_1 suspend boom!")); @@ -3232,8 +2994,6 @@ public void suspend() { assertThat(task00.state(), is(Task.State.CREATED)); assertThat(task01.state(), is(Task.State.SUSPENDED)); assertThat(task02.state(), is(Task.State.SUSPENDED)); - - verify(activeTaskCreator); } @Test @@ -3265,16 +3025,11 @@ public void suspend() { } }; - resetToStrict(changeLogReader); - changeLogReader.enforceRestoreActive(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(asList(task00, task01, task02)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(anyObject()); - expectLastCall().andThrow(new RuntimeException("whatever")).times(3); - activeTaskCreator.closeThreadProducerIfNeeded(); - expectLastCall().andThrow(new RuntimeException("whatever all")); - expect(standbyTaskCreator.createTasks(eq(emptyMap()))).andStubReturn(emptyList()); - replay(activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01, task02)); + doThrow(new RuntimeException("whatever")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(any()); + doThrow(new RuntimeException("whatever all")).when(activeTaskCreator).closeThreadProducerIfNeeded(); + when(standbyTaskCreator.createTasks(eq(emptyMap()))).thenReturn(emptyList()); taskManager.handleAssignment(assignment, emptyMap()); @@ -3306,8 +3061,6 @@ public void suspend() { assertThat(task02.state(), is(Task.State.CLOSED)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - // the active task creator should also get closed (so that it closes the thread producer if applicable) - verify(activeTaskCreator, changeLogReader); } @Test @@ -3316,22 +3069,12 @@ public void shouldCloseStandbyTasksOnShutdown() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false); // `handleAssignment` - expect(activeTaskCreator.createTasks(anyObject(), anyObject())).andStubReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(assignment))).andStubReturn(singletonList(task00)); + when(activeTaskCreator.createTasks(any(), any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(assignment))).thenReturn(singletonList(task00)); // `tryToCompleteRestoration` - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); - - // `shutdown` - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - activeTaskCreator.closeThreadProducerIfNeeded(); - expectLastCall(); - - replay(consumer, activeTaskCreator, standbyTaskCreator, changeLogReader); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(consumer.assignment()).thenReturn(emptySet()); taskManager.handleAssignment(emptyMap(), assignment); assertThat(task00.state(), is(Task.State.CREATED)); @@ -3345,8 +3088,9 @@ public void shouldCloseStandbyTasksOnShutdown() { assertThat(task00.state(), is(Task.State.CLOSED)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - // the active task creator should also get closed (so that it closes the thread producer if applicable) - verify(activeTaskCreator); + verify(consumer).resume(eq(emptySet())); + // `shutdown` + verify(activeTaskCreator).closeThreadProducerIfNeeded(); } @Test @@ -3361,18 +3105,16 @@ public void shouldShutDownStateUpdaterAndCloseFailedTasksDirty() { new ExceptionAndTasks(mkSet(failedStatefulTask), new RuntimeException()), new ExceptionAndTasks(mkSet(failedStandbyTask), new RuntimeException())) ); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(failedStatefulTask.id()); - activeTaskCreator.closeThreadProducerIfNeeded(); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(activeTaskCreator); taskManager.shutdown(true); - verify(activeTaskCreator); - Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); - Mockito.verify(failedStatefulTask).prepareCommit(); - Mockito.verify(failedStatefulTask).suspend(); - Mockito.verify(failedStatefulTask).closeDirty(); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(failedStatefulTask.id()); + verify(activeTaskCreator).closeThreadProducerIfNeeded(); + verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); + verify(failedStatefulTask).prepareCommit(); + verify(failedStatefulTask).suspend(); + verify(failedStatefulTask).closeDirty(); } @Test @@ -3386,19 +3128,17 @@ public void shouldShutDownStateUpdaterAndAddRestoredTasksToTaskRegistry() { final Set restoredTasks = restoredActiveTasks.stream().map(t -> (Task) t).collect(Collectors.toSet()); when(stateUpdater.drainRestoredActiveTasks(Duration.ZERO)).thenReturn(restoredActiveTasks); when(tasks.activeTasks()).thenReturn(restoredTasks); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask1.id()); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(statefulTask2.id()); - activeTaskCreator.closeThreadProducerIfNeeded(); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(activeTaskCreator); taskManager.shutdown(true); - verify(activeTaskCreator); - Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); - Mockito.verify(tasks).addActiveTasks(restoredTasks); - Mockito.verify(statefulTask1).closeClean(); - Mockito.verify(statefulTask2).closeClean(); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask1.id()); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(statefulTask2.id()); + verify(activeTaskCreator).closeThreadProducerIfNeeded(); + verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); + verify(tasks).addActiveTasks(restoredTasks); + verify(statefulTask1).closeClean(); + verify(statefulTask2).closeClean(); } @Test @@ -3411,29 +3151,26 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() { when(stateUpdater.drainRemovedTasks()).thenReturn(mkSet(removedStandbyTask, removedStatefulTask)); when(tasks.activeTasks()).thenReturn(mkSet(removedStatefulTask)); when(tasks.allTasks()).thenReturn(mkSet(removedStatefulTask, removedStandbyTask)); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(removedStatefulTask.id()); - activeTaskCreator.closeThreadProducerIfNeeded(); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(activeTaskCreator); taskManager.shutdown(true); - verify(activeTaskCreator); - Mockito.verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); - Mockito.verify(tasks).addActiveTasks(mkSet(removedStatefulTask)); - Mockito.verify(tasks).addStandbyTasks(mkSet(removedStandbyTask)); - Mockito.verify(removedStatefulTask).closeClean(); - Mockito.verify(removedStandbyTask).closeClean(); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(removedStatefulTask.id()); + verify(activeTaskCreator).closeThreadProducerIfNeeded(); + verify(stateUpdater).shutdown(Duration.ofMillis(Long.MAX_VALUE)); + verify(tasks).addActiveTasks(mkSet(removedStatefulTask)); + verify(tasks).addStandbyTasks(mkSet(removedStandbyTask)); + verify(removedStatefulTask).closeClean(); + verify(removedStandbyTask).closeClean(); } @Test public void shouldInitializeNewActiveTasks() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))) - .andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))) + .thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3441,8 +3178,7 @@ public void shouldInitializeNewActiveTasks() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - // verifies that we actually resume the assignment at the end of restoration. - verify(consumer); + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3450,10 +3186,8 @@ public void shouldInitializeNewStandbyTasks() { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), anyObject())).andStubReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))).andStubReturn(singletonList(task01)); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))).thenReturn(singletonList(task01)); taskManager.handleAssignment(emptyMap(), taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3461,21 +3195,21 @@ public void shouldInitializeNewStandbyTasks() { assertThat(task01.state(), is(Task.State.RUNNING)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.equalTo(singletonMap(taskId01, task01))); + + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldHandleRebalanceEvents() { final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.pause(assignment); - expectLastCall(); - expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()); - replay(consumer, stateDirectory); + when(consumer.assignment()).thenReturn(assignment); + when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>()); assertThat(taskManager.rebalanceInProgress(), is(false)); taskManager.handleRebalanceStart(emptySet()); assertThat(taskManager.rebalanceInProgress(), is(true)); taskManager.handleRebalanceComplete(); assertThat(taskManager.rebalanceInProgress(), is(false)); + verify(consumer).pause(assignment); } @Test @@ -3486,14 +3220,10 @@ public void shouldCommitActiveAndStandbyTasks() { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))) - .andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))) - .andStubReturn(singletonList(task01)); - consumer.commitSync(offsets); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))) + .thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))) + .thenReturn(singletonList(task01)); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3507,6 +3237,9 @@ public void shouldCommitActiveAndStandbyTasks() { assertThat(taskManager.commitAll(), equalTo(2)); assertThat(task00.commitNeeded, is(false)); assertThat(task01.commitNeeded, is(false)); + verify(consumer).commitSync(offsets); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3530,15 +3263,13 @@ public void shouldCommitProvidedTasksIfNeeded() { ); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))) - .andStubReturn(Arrays.asList(task00, task01, task02)); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))) - .andStubReturn(Arrays.asList(task03, task04, task05)); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))) + .thenReturn(Arrays.asList(task00, task01, task02)); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))) + .thenReturn(Arrays.asList(task03, task04, task05)); consumer.commitSync(eq(emptyMap())); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); - taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3557,6 +3288,8 @@ public void shouldCommitProvidedTasksIfNeeded() { assertThat(task03.commitNeeded, is(false)); assertThat(task04.commitNeeded, is(true)); assertThat(task05.commitNeeded, is(false)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3564,11 +3297,8 @@ public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, false); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), anyObject())).andStubReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(taskId00Assignment))).thenReturn(singletonList(task00)); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3579,6 +3309,8 @@ public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { assertThat(taskManager.commitAll(), equalTo(1)); assertThat(task00.commitNeeded, is(false)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3589,12 +3321,10 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw makeTaskFolders(taskId00.toString(), task01.toString()); expectLockObtainedFor(taskId00, taskId01); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))) - .andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))) - .andStubReturn(singletonList(task01)); - - replay(activeTaskCreator, standbyTaskCreator, stateDirectory, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))) + .thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))) + .thenReturn(singletonList(task01)); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3616,6 +3346,8 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(-1) // sentinel indicating that nothing was done because a rebalance is in progress ); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3626,37 +3358,28 @@ public void shouldCommitViaConsumerIfEosDisabled() { task01.setCommitNeeded(); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall(); - replay(consumer); - taskManager.commitAll(); - verify(consumer); + verify(consumer).commitSync(offsets); } @Test public void shouldCommitViaProducerIfEosAlphaEnabled() { - final StreamsProducer producer = EasyMock.mock(StreamsProducer.class); - expect(activeTaskCreator.streamsProducerForTask(anyObject(TaskId.class))) - .andReturn(producer) - .andReturn(producer); + final StreamsProducer producer = mock(StreamsProducer.class); + when(activeTaskCreator.streamsProducerForTask(any(TaskId.class))).thenReturn(producer); final Map offsetsT01 = singletonMap(t1p1, new OffsetAndMetadata(0L, null)); final Map offsetsT02 = singletonMap(t1p2, new OffsetAndMetadata(1L, null)); - producer.commitTransaction(offsetsT01, new ConsumerGroupMetadata("appId")); - expectLastCall(); - producer.commitTransaction(offsetsT02, new ConsumerGroupMetadata("appId")); - expectLastCall(); - shouldCommitViaProducerIfEosEnabled(ProcessingMode.EXACTLY_ONCE_ALPHA, producer, offsetsT01, offsetsT02); + verify(producer).commitTransaction(offsetsT01, new ConsumerGroupMetadata("appId")); + verify(producer).commitTransaction(offsetsT02, new ConsumerGroupMetadata("appId")); } @Test public void shouldCommitViaProducerIfEosV2Enabled() { - final StreamsProducer producer = EasyMock.mock(StreamsProducer.class); - expect(activeTaskCreator.threadProducer()).andReturn(producer); + final StreamsProducer producer = mock(StreamsProducer.class); + when(activeTaskCreator.threadProducer()).thenReturn(producer); final Map offsetsT01 = singletonMap(t1p1, new OffsetAndMetadata(0L, null)); final Map offsetsT02 = singletonMap(t1p2, new OffsetAndMetadata(1L, null)); @@ -3664,10 +3387,9 @@ public void shouldCommitViaProducerIfEosV2Enabled() { allOffsets.putAll(offsetsT01); allOffsets.putAll(offsetsT02); - producer.commitTransaction(allOffsets, new ConsumerGroupMetadata("appId")); - expectLastCall(); - shouldCommitViaProducerIfEosEnabled(ProcessingMode.EXACTLY_ONCE_V2, producer, offsetsT01, offsetsT02); + + verify(producer).commitTransaction(allOffsets, new ConsumerGroupMetadata("appId")); } private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processingMode, @@ -3685,13 +3407,9 @@ private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processing task02.setCommitNeeded(); taskManager.addTask(task02); - reset(consumer); - expect(consumer.groupMetadata()).andStubReturn(new ConsumerGroupMetadata("appId")); - replay(activeTaskCreator, consumer, producer); + when(consumer.groupMetadata()).thenReturn(new ConsumerGroupMetadata("appId")); taskManager.commitAll(); - - verify(producer, consumer); } @Test @@ -3704,9 +3422,8 @@ public Map prepareCommit() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3718,6 +3435,8 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3730,10 +3449,8 @@ public Map prepareCommit() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), anyObject())).andStubReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(taskId01Assignment))).andStubReturn(singletonList(task01)); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(taskId01Assignment))).thenReturn(singletonList(task01)); taskManager.handleAssignment(emptyMap(), taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3745,16 +3462,16 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); + + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldSendPurgeData() { - resetToStrict(adminClient); - expect(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L)))) - .andReturn(new DeleteRecordsResult(singletonMap(t1p1, completedFuture()))); - expect(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(17L)))) - .andReturn(new DeleteRecordsResult(singletonMap(t1p1, completedFuture()))); - replay(adminClient); + when(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L)))) + .thenReturn(new DeleteRecordsResult(singletonMap(t1p1, completedFuture()))); + when(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(17L)))) + .thenReturn(new DeleteRecordsResult(singletonMap(t1p1, completedFuture()))); final Map purgableOffsets = new HashMap<>(); final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true) { @@ -3765,10 +3482,8 @@ public Map purgeableOffsets() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3781,16 +3496,14 @@ public Map purgeableOffsets() { purgableOffsets.put(t1p1, 17L); taskManager.maybePurgeCommittedRecords(); - verify(adminClient); + verifyConsumerResumedWithAssignment(consumer); } @Test public void shouldNotSendPurgeDataIfPreviousNotDone() { - resetToStrict(adminClient); final KafkaFutureImpl futureDeletedRecords = new KafkaFutureImpl<>(); - expect(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L)))) - .andReturn(new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords))); - replay(adminClient); + when(adminClient.deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L)))) + .thenReturn(new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords))); final Map purgableOffsets = new HashMap<>(); final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true) { @@ -3801,9 +3514,8 @@ public Map purgeableOffsets() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3819,7 +3531,7 @@ public Map purgeableOffsets() { purgableOffsets.put(t1p1, 17L); taskManager.maybePurgeCommittedRecords(); - verify(adminClient); + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3831,9 +3543,7 @@ public void shouldIgnorePurgeDataErrors() { final KafkaFutureImpl futureDeletedRecords = new KafkaFutureImpl<>(); final DeleteRecordsResult deleteRecordsResult = new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords)); futureDeletedRecords.completeExceptionally(new Exception("KABOOM!")); - expect(adminClient.deleteRecords(anyObject())).andReturn(deleteRecordsResult).times(2); - - replay(activeTaskCreator, adminClient, consumer, changeLogReader); + when(adminClient.deleteRecords(any())).thenReturn(deleteRecordsResult); taskManager.addTask(task00); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3845,7 +3555,7 @@ public void shouldIgnorePurgeDataErrors() { taskManager.maybePurgeCommittedRecords(); taskManager.maybePurgeCommittedRecords(); - verify(adminClient); + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3878,14 +3588,10 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { ); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignmentActive))) - .andStubReturn(asList(task00, task01, task02, task03)); - expect(standbyTaskCreator.createTasks(eq(assignmentStandby))) - .andStubReturn(singletonList(task04)); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignmentActive))) + .thenReturn(asList(task00, task01, task02, task03)); + when(standbyTaskCreator.createTasks(eq(assignmentStandby))) + .thenReturn(singletonList(task04)); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3910,6 +3616,9 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { task04.setCommitRequested(); assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3)); + verify(consumer).commitSync(expectedCommittedOffsets); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -3922,9 +3631,8 @@ public void shouldProcessActiveTasks() { assignment.put(taskId01, taskId01Partitions); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andStubReturn(Arrays.asList(task00, task01)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(Arrays.asList(task00, task01)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3960,6 +3668,8 @@ public void shouldProcessActiveTasks() { // check that if there's no records proccssible, we would stop early assertThat(taskManager.process(3, time), is(5)); assertThat(taskManager.process(3, time), is(0)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4035,9 +3745,8 @@ public boolean process(final long wallClockTime) { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4048,6 +3757,8 @@ public boolean process(final long wallClockTime) { task00.addRecords(partition, singletonList(getConsumerRecord(partition, 0L))); assertThrows(TaskMigratedException.class, () -> taskManager.process(1, time)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4060,10 +3771,9 @@ public boolean process(final long wallClockTime) { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))) - .andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))) + .thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4077,6 +3787,8 @@ public boolean process(final long wallClockTime) { assertThat(exception.taskId().isPresent(), is(true)); assertThat(exception.taskId().get(), is(taskId00)); assertThat(exception.getCause().getMessage(), is("oops")); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4089,9 +3801,8 @@ public boolean maybePunctuateStreamTime() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4099,6 +3810,8 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(TaskMigratedException.class, () -> taskManager.punctuate()); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4111,9 +3824,8 @@ public boolean maybePunctuateStreamTime() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4121,6 +3833,8 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(KafkaException.class, () -> taskManager.punctuate()); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4138,11 +3852,10 @@ public boolean maybePunctuateSystemTime() { }; expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))) - .andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())) - .andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))) + .thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())) + .thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4151,6 +3864,8 @@ public boolean maybePunctuateSystemTime() { // one for stream and one for system time assertThat(taskManager.punctuate(), equalTo(2)); + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4162,17 +3877,16 @@ public Set changelogPartitions() { } }; - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, changeLogReader, consumer); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(false)); assertThat(task00.state(), is(Task.State.RESTORING)); // this could be a bit mysterious; we're verifying _no_ interactions on the consumer, // since the taskManager should _not_ resume the assignment while we're still in RESTORING - verify(consumer); + verifyNoInteractions(consumer); } @Test @@ -4182,12 +3896,8 @@ public void shouldHaveRemainingPartitionsUncleared() { task00.setCommittableOffsetsAndMetadata(offsets); expectRestoreToBeCompleted(consumer, changeLogReader); - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(task00)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - consumer.commitSync(offsets); - expectLastCall(); - - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(task00)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) { LogCaptureAppender.setClassLoggerToDebug(TaskManager.class); @@ -4207,6 +3917,8 @@ public void shouldHaveRemainingPartitionsUncleared() { "tasks have been cleaned up by the handleAssignment callback.") ); } + + verifyConsumerResumedWithAssignment(consumer); } @Test @@ -4314,8 +4026,7 @@ public void shouldTransmitProducerMetrics() { new MockTime()); final Map dummyProducerMetrics = singletonMap(testMetricName, testMetric); - expect(activeTaskCreator.producerMetrics()).andReturn(dummyProducerMetrics); - replay(activeTaskCreator); + when(activeTaskCreator.producerMetrics()).thenReturn(dummyProducerMetrics); assertThat(taskManager.producerMetrics(), is(dummyProducerMetrics)); } @@ -4341,11 +4052,10 @@ private Map handleAssignment(final Map allActiveTasks = new HashSet<>(runningTasks); allActiveTasks.addAll(restoringTasks); - expect(standbyTaskCreator.createTasks(eq(standbyAssignment))).andStubReturn(standbyTasks); - expect(activeTaskCreator.createTasks(anyObject(), eq(allActiveTasksAssignment))).andStubReturn(allActiveTasks); + when(standbyTaskCreator.createTasks(eq(standbyAssignment))).thenReturn(standbyTasks); + when(activeTaskCreator.createTasks(any(), eq(allActiveTasksAssignment))).thenReturn(allActiveTasks); expectRestoreToBeCompleted(consumer, changeLogReader); - replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader); taskManager.handleAssignment(allActiveTasksAssignment, standbyAssignment); taskManager.tryToCompleteRestoration(time.milliseconds(), null); @@ -4370,27 +4080,26 @@ private Map handleAssignment(final Map consumer) { final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.pause(assignment); + when(consumer.assignment()).thenReturn(assignment); + doNothing().when(consumer).pause(assignment); } @Test @@ -4401,9 +4110,7 @@ public void shouldThrowTaskMigratedExceptionOnCommitFailed() { task01.setCommitNeeded(); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new CommitFailedException()); - replay(consumer); + doThrow(new CommitFailedException()).when(consumer).commitSync(offsets); final TaskMigratedException thrown = assertThrows( TaskMigratedException.class, @@ -4428,11 +4135,7 @@ public void shouldNotFailForTimeoutExceptionOnConsumerCommit() { task00.setCommittableOffsetsAndMetadata(taskId00Partitions.stream().collect(Collectors.toMap(p -> p, p -> new OffsetAndMetadata(0)))); task01.setCommittableOffsetsAndMetadata(taskId00Partitions.stream().collect(Collectors.toMap(p -> p, p -> new OffsetAndMetadata(0)))); - consumer.commitSync(anyObject(Map.class)); - expectLastCall().andThrow(new TimeoutException("KABOOM!")); - consumer.commitSync(anyObject(Map.class)); - expectLastCall(); - replay(consumer); + doThrow(new TimeoutException("KABOOM!")).doNothing().when(consumer).commitSync(any(Map.class)); task00.setCommitNeeded(); @@ -4451,10 +4154,7 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() { final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_ALPHA, tasks, false); final StreamsProducer producer = mock(StreamsProducer.class); - expect(activeTaskCreator.streamsProducerForTask(anyObject(TaskId.class))) - .andReturn(producer) - .andReturn(producer) - .andReturn(producer); + when(activeTaskCreator.streamsProducerForTask(any(TaskId.class))).thenReturn(producer); final Map offsetsT00 = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); final Map offsetsT01 = singletonMap(t1p1, new OffsetAndMetadata(1L, null)); @@ -4475,8 +4175,7 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() { final StateMachineTask task02 = new StateMachineTask(taskId02, taskId02Partitions, true); when(tasks.allTasks()).thenReturn(mkSet(task00, task01, task02)); - expect(consumer.groupMetadata()).andStubReturn(null); - replay(activeTaskCreator, consumer); + when(consumer.groupMetadata()).thenReturn(null); task00.setCommitNeeded(); task01.setCommitNeeded(); @@ -4496,9 +4195,7 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false); final StreamsProducer producer = mock(StreamsProducer.class); - expect(activeTaskCreator.threadProducer()) - .andReturn(producer) - .andReturn(producer); + when(activeTaskCreator.threadProducer()).thenReturn(producer); final Map offsetsT00 = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); final Map offsetsT01 = singletonMap(t1p1, new OffsetAndMetadata(1L, null)); @@ -4513,8 +4210,7 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV task01.setCommittableOffsetsAndMetadata(offsetsT01); final StateMachineTask task02 = new StateMachineTask(taskId02, taskId02Partitions, true); - expect(consumer.groupMetadata()).andStubReturn(null); - replay(activeTaskCreator, consumer); + when(consumer.groupMetadata()).thenReturn(null); task00.setCommitNeeded(); task01.setCommitNeeded(); @@ -4537,9 +4233,7 @@ public void shouldStreamsExceptionOnCommitError() { task01.setCommitNeeded(); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new KafkaException()); - replay(consumer); + doThrow(new KafkaException()).when(consumer).commitSync(offsets); final StreamsException thrown = assertThrows( StreamsException.class, @@ -4559,9 +4253,7 @@ public void shouldFailOnCommitFatal() { task01.setCommitNeeded(); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new RuntimeException("KABOOM")); - replay(consumer); + doThrow(new RuntimeException("KABOOM")).when(consumer).commitSync(offsets); final RuntimeException thrown = assertThrows( RuntimeException.class, @@ -4585,9 +4277,8 @@ public void suspend() { final Map> assignment = new HashMap<>(taskId00Assignment); assignment.putAll(taskId01Assignment); - expect(activeTaskCreator.createTasks(anyObject(), eq(assignment))).andReturn(asList(task00, task01)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - replay(activeTaskCreator, standbyTaskCreator, consumer); + when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); taskManager.handleAssignment(assignment, Collections.emptyMap()); @@ -4602,30 +4293,24 @@ public void suspend() { @Test public void shouldConvertActiveTaskToStandbyTask() { - final StreamTask activeTask = EasyMock.mock(StreamTask.class); - expect(activeTask.id()).andStubReturn(taskId00); - expect(activeTask.inputPartitions()).andStubReturn(taskId00Partitions); - expect(activeTask.isActive()).andStubReturn(true); - expect(activeTask.prepareCommit()).andStubReturn(Collections.emptyMap()); - - final StandbyTask standbyTask = EasyMock.mock(StandbyTask.class); - expect(standbyTask.id()).andStubReturn(taskId00); - - expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andReturn(singletonList(activeTask)); - expect(standbyTaskCreator.createTasks(anyObject())).andStubReturn(Collections.emptySet()); - activeTask.prepareRecycle(); - expectLastCall().once(); - activeTaskCreator.closeAndRemoveTaskProducerIfNeeded(taskId00); - expectLastCall().anyTimes(); - expect(standbyTaskCreator.createStandbyTaskFromActive(anyObject(), eq(taskId00Partitions))).andReturn(standbyTask); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); + final StreamTask activeTask = mock(StreamTask.class); + when(activeTask.id()).thenReturn(taskId00); + when(activeTask.inputPartitions()).thenReturn(taskId00Partitions); + when(activeTask.isActive()).thenReturn(true); + when(activeTask.prepareCommit()).thenReturn(Collections.emptyMap()); + + final StandbyTask standbyTask = mock(StandbyTask.class); + when(standbyTask.id()).thenReturn(taskId00); - replay(activeTask, standbyTask, activeTaskCreator, standbyTaskCreator, consumer); + when(activeTaskCreator.createTasks(any(), eq(taskId00Assignment))).thenReturn(singletonList(activeTask)); + when(standbyTaskCreator.createTasks(any())).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createStandbyTaskFromActive(any(), eq(taskId00Partitions))).thenReturn(standbyTask); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap()); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); - verify(activeTaskCreator, standbyTaskCreator); + verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); } @Test @@ -4639,18 +4324,14 @@ public void shouldConvertStandbyTaskToActiveTask() { when(activeTask.id()).thenReturn(taskId00); when(activeTask.inputPartitions()).thenReturn(taskId00Partitions); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(taskId00Assignment))).andReturn(singletonList(standbyTask)); - expect(activeTaskCreator.createActiveTaskFromStandby(eq(standbyTask), eq(taskId00Partitions), anyObject())).andReturn(activeTask); - expect(activeTaskCreator.createTasks(anyObject(), eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - expect(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).andReturn(Collections.emptySet()); - - replay(standbyTaskCreator, activeTaskCreator, consumer); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(taskId00Assignment))).thenReturn(singletonList(standbyTask)); + when(activeTaskCreator.createActiveTaskFromStandby(eq(standbyTask), eq(taskId00Partitions), any())).thenReturn(activeTask); + when(activeTaskCreator.createTasks(any(), eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); + when(standbyTaskCreator.createTasks(eq(Collections.emptyMap()))).thenReturn(Collections.emptySet()); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap()); - - verify(standbyTaskCreator, activeTaskCreator); } @Test @@ -4666,17 +4347,14 @@ public void shouldListNotPausedTasks() { private static void expectRestoreToBeCompleted(final Consumer consumer, final ChangelogReader changeLogReader) { - expectRestoreToBeCompleted(consumer, changeLogReader, true); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(consumer.assignment()).thenReturn(assignment); + when(changeLogReader.completedChangelogs()).thenReturn(emptySet()); } - private static void expectRestoreToBeCompleted(final Consumer consumer, - final ChangelogReader changeLogReader, - final boolean changeLogUpdateRequired) { + private static void verifyConsumerResumedWithAssignment(final Consumer consumer) { final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.resume(assignment); - expectLastCall(); - expect(changeLogReader.completedChangelogs()).andReturn(emptySet()).times(changeLogUpdateRequired ? 1 : 0, 1); + verify(consumer, atLeastOnce()).resume(assignment); } private static KafkaFutureImpl completedFuture() { @@ -4690,14 +4368,14 @@ private void makeTaskFolders(final String... names) throws Exception { for (int i = 0; i < names.length; ++i) { taskFolders.add(new TaskDirectory(testFolder.newFolder(names[i]), null)); } - expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(taskFolders).once(); + when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(taskFolders).thenReturn(Collections.emptyList()); } private void writeCheckpointFile(final TaskId task, final Map offsets) throws Exception { final File checkpointFile = getCheckpointFile(task); Files.createFile(checkpointFile.toPath()); new OffsetCheckpoint(checkpointFile).write(offsets); - expect(stateDirectory.checkpointFileFor(task)).andReturn(checkpointFile); + when(stateDirectory.checkpointFileFor(task)).thenReturn(checkpointFile); } private File getCheckpointFile(final TaskId task) {