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 61e078868bf71..7d048cb973d2b 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 @@ -108,7 +108,6 @@ import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask; -import static org.easymock.EasyMock.anyObject; import static org.easymock.EasyMock.eq; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.expectLastCall; @@ -198,6 +197,8 @@ public class TaskManagerTest { @Mock(type = MockType.STRICT) private Consumer consumer; @org.mockito.Mock + private Consumer mockitoConsumer; + @org.mockito.Mock private ActiveTaskCreator activeTaskCreator; @org.mockito.Mock private StandbyTaskCreator standbyTaskCreator; @@ -310,16 +311,15 @@ public void shouldLockAllTasksOnCorruptionWithProcessingThreads() { .withInputPartitions(taskId00Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true, true); + taskManager.setMainConsumer(mockitoConsumer); when(tasks.activeTaskIds()).thenReturn(mkSet(taskId00, taskId01)); when(tasks.task(taskId00)).thenReturn(activeTask1); final KafkaFuture mockFuture = KafkaFuture.completedFuture(null); when(schedulingTaskManager.lockTasks(any())).thenReturn(mockFuture); - expect(consumer.assignment()).andReturn(emptySet()).anyTimes(); - replay(consumer); taskManager.handleCorruption(mkSet(taskId00)); - verify(consumer); + Mockito.verify(mockitoConsumer).assignment(); Mockito.verify(schedulingTaskManager).lockTasks(mkSet(taskId00, taskId01)); Mockito.verify(schedulingTaskManager).unlockTasks(mkSet(taskId00, taskId01)); } @@ -1210,13 +1210,13 @@ public void shouldSuspendRevokedTaskRemovedFromStateUpdater() { when(stateUpdater.hasRemovedTasks()).thenReturn(true); when(stateUpdater.drainRemovedTasks()).thenReturn(mkSet(statefulTask)); taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); Mockito.verify(statefulTask).suspend(); Mockito.verify(tasks).addTask(statefulTask); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -1242,13 +1242,10 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() { when(stateUpdater.drainRemovedTasks()) .thenReturn(mkSet(taskToRecycle0, taskToRecycle1, taskToClose, taskToUpdateInputPartitions, taskToCloseReviveAndUpdateInputPartitions)); when(stateUpdater.restoresActiveTasks()).thenReturn(true); - when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, consumer)) + when(activeTaskCreator.createActiveTaskFromStandby(taskToRecycle1, taskId01Partitions, mockitoConsumer)) .thenReturn(convertedTask1); when(standbyTaskCreator.createStandbyTaskFromActive(taskToRecycle0, taskId00Partitions)) .thenReturn(convertedTask0); - expect(consumer.assignment()).andReturn(emptySet()).anyTimes(); - consumer.resume(anyObject()); - expectLastCall().anyTimes(); final TasksRegistry tasks = mock(TasksRegistry.class); when(tasks.removePendingTaskToCloseClean(taskToClose.id())).thenReturn(true); when(tasks.removePendingTaskToCloseClean(argThat(taskId -> !taskId.equals(taskToClose.id())))).thenReturn(false); @@ -1263,12 +1260,11 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() { argThat(taskId -> !taskId.equals(taskToCloseReviveAndUpdateInputPartitions.id())) )).thenReturn(null); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - taskManager.setMainConsumer(consumer); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter -> { }); - verify(consumer); + Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any()); Mockito.verify(activeTaskCreator, times(3)).closeAndRemoveTaskProducerIfNeeded(any()); Mockito.verify(convertedTask0).initializeIfNeeded(); Mockito.verify(convertedTask1).initializeIfNeeded(); @@ -1282,6 +1278,7 @@ public void shouldHandleMultipleRemovedTasksFromStateUpdater() { Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).updateInputPartitions(Mockito.eq(taskId05Partitions), anyMap()); Mockito.verify(taskToCloseReviveAndUpdateInputPartitions).initializeIfNeeded(); Mockito.verify(stateUpdater).add(taskToCloseReviveAndUpdateInputPartitions); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -1417,15 +1414,14 @@ public void shouldTransitRestoredTaskToRunning() { .withInputPartitions(taskId00Partitions).build(); final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks); - consumer.resume(task.inputPartitions()); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); Mockito.verify(task).completeRestoration(noOpResetter); Mockito.verify(task).clearTaskTimeout(); Mockito.verify(tasks).addTask(task); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(task.inputPartitions()); } @Test @@ -1435,16 +1431,16 @@ public void shouldHandleTimeoutExceptionInTransitRestoredTaskToRunning() { .withInputPartitions(taskId00Partitions).build(); final TasksRegistry tasks = mock(TasksRegistry.class); final TaskManager taskManager = setUpTransitionToRunningOfRestoredTask(task, tasks); + taskManager.setMainConsumer(mockitoConsumer); 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); + Mockito.verifyNoInteractions(mockitoConsumer); } private TaskManager setUpTransitionToRunningOfRestoredTask(final StreamTask statefulTask, @@ -1634,12 +1630,11 @@ public void shouldUpdateInputPartitionsOfRestoredTask() { when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(mkSet(statefulTask)); when(stateUpdater.restoresActiveTasks()).thenReturn(true); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - consumer.resume(statefulTask.inputPartitions()); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(statefulTask.inputPartitions()); Mockito.verify(statefulTask).updateInputPartitions(Mockito.eq(taskId01Partitions), anyMap()); Mockito.verify(statefulTask).completeRestoration(noOpResetter); Mockito.verify(statefulTask).clearTaskTimeout(); @@ -1680,13 +1675,13 @@ 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.setMainConsumer(mockitoConsumer); taskManager.checkStateUpdater(time.milliseconds(), noOpResetter); - verify(consumer); Mockito.verify(statefulTask).suspend(); Mockito.verify(tasks).addTask(statefulTask); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -1937,13 +1932,12 @@ public void shouldUnlockEmptyDirsAtRebalanceStart() throws Exception { @Test public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() { final Set assigned = mkSet(t1p0, t1p1); - expect(consumer.assignment()).andReturn(assigned); - consumer.pause(assigned); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); + when(mockitoConsumer.assignment()).thenReturn(assigned); taskManager.handleRebalanceComplete(); - verify(consumer); + Mockito.verify(mockitoConsumer).pause(assigned); } @Test @@ -1953,15 +1947,14 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() { .withInputPartitions(taskId00Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); + taskManager.setMainConsumer(mockitoConsumer); when(tasks.allTasks()).thenReturn(mkSet(statefulTask0)); final Set assigned = mkSet(t1p0, t1p1); - expect(consumer.assignment()).andReturn(assigned); - consumer.pause(mkSet(t1p1)); - replay(consumer); + when(mockitoConsumer.assignment()).thenReturn(assigned); taskManager.handleRebalanceComplete(); - verify(consumer); + Mockito.verify(mockitoConsumer).pause(mkSet(t1p1)); } @Test @@ -2006,6 +1999,7 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater() .withInputPartitions(taskId03Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); + taskManager.setMainConsumer(mockitoConsumer); when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId00, runningStatefulTask))); when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask, restoringStatefulTask)); when(tasks.allTasks()).thenReturn(mkSet(runningStatefulTask)); @@ -2021,14 +2015,12 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater() replay(stateDirectory); final Set assigned = mkSet(t1p0, t1p1, t1p2); - expect(consumer.assignment()).andReturn(assigned); - consumer.pause(mkSet(t1p1, t1p2)); - replay(consumer); + when(mockitoConsumer.assignment()).thenReturn(assigned); taskManager.handleRebalanceStart(singleton("topic")); taskManager.handleRebalanceComplete(); - verify(consumer); + Mockito.verify(mockitoConsumer).pause(mkSet(t1p1, t1p2)); verify(stateDirectory); assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02))); } @@ -2460,10 +2452,9 @@ public void shouldReAddRevivedTasksToStateUpdater() { .withInputPartitions(taskId02Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); + taskManager.setMainConsumer(mockitoConsumer); when(tasks.task(taskId03)).thenReturn(corruptedActiveTask); when(tasks.task(taskId02)).thenReturn(corruptedStandbyTask); - expect(consumer.assignment()).andReturn(emptySet()); - replay(consumer); taskManager.handleCorruption(mkSet(corruptedActiveTask.id(), corruptedStandbyTask.id())); @@ -2477,6 +2468,7 @@ public void shouldReAddRevivedTasksToStateUpdater() { Mockito.verify(tasks).removeTask(corruptedStandbyTask); Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedActiveTask)); Mockito.verify(tasks).addPendingTasksToInit(mkSet(corruptedStandbyTask)); + Mockito.verify(mockitoConsumer).assignment(); } @Test @@ -2635,9 +2627,9 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndNotCommitRunningSt when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId02, corruptedTask))); when(tasks.task(taskId02)).thenReturn(corruptedTask); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); + taskManager.setMainConsumer(mockitoConsumer); when(stateUpdater.getTasks()).thenReturn(mkSet(activeRestoringTask, standbyTask)); - expect(consumer.assignment()).andReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - replay(consumer); + when(mockitoConsumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); taskManager.handleCorruption(mkSet(taskId02)); @@ -2647,7 +2639,6 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndNotCommitRunningSt Mockito.verify(standbyTask, never()).commitNeeded(); Mockito.verify(standbyTask, never()).prepareCommit(); Mockito.verify(standbyTask, never()).postCommit(Mockito.anyBoolean()); - verify(consumer); } @Test @@ -2670,8 +2661,8 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndCommitRunningStand )); when(tasks.task(taskId02)).thenReturn(corruptedTask); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, false); - expect(consumer.assignment()).andReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); + when(mockitoConsumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); taskManager.handleCorruption(mkSet(taskId02)); @@ -3114,11 +3105,9 @@ public void shouldAddNewActiveTasks() { final Map> assignment = taskId00Assignment; final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); + taskManager.setMainConsumer(mockitoConsumer); + when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(singletonList(task00)); - replay(consumer); taskManager.handleAssignment(assignment, emptyMap()); @@ -3130,6 +3119,8 @@ public void shouldAddNewActiveTasks() { assertThat(taskManager.activeTaskMap(), Matchers.equalTo(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); Mockito.verify(changeLogReader).enforceRestoreActive(); + Mockito.verify(mockitoConsumer).assignment(); + Mockito.verify(mockitoConsumer).resume(Mockito.eq(emptySet())); } @Test @@ -3151,13 +3142,9 @@ public void initializeIfNeeded() { } }; - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); + taskManager.setMainConsumer(mockitoConsumer); + when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01)); - replay(consumer); taskManager.handleAssignment(assignment, emptyMap()); @@ -3174,6 +3161,7 @@ public void initializeIfNeeded() { ); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); Mockito.verify(changeLogReader).enforceRestoreActive(); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -3188,14 +3176,9 @@ public void completeRestoration(final java.util.function.Consumer> assignment = singletonMap(taskId00, taskId00Partitions); final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); + taskManager.setMainConsumer(mockitoConsumer); + // `handleAssignment` when(standbyTaskCreator.createTasks(assignment)).thenReturn(singletonList(task00)); - // `tryToCompleteRestoration` - expect(consumer.assignment()).andReturn(emptySet()); - consumer.resume(eq(emptySet())); - expectLastCall(); - - // `shutdown` - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - - replay(consumer); - taskManager.handleAssignment(emptyMap(), assignment); assertThat(task00.state(), is(Task.State.CREATED)); @@ -3833,6 +3807,9 @@ public void shouldCloseStandbyTasksOnShutdown() { assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); // the active task creator should also get closed (so that it closes the thread producer if applicable) Mockito.verify(activeTaskCreator).closeThreadProducerIfNeeded(); + // `tryToCompleteRestoration` + Mockito.verify(mockitoConsumer).assignment(); + Mockito.verify(mockitoConsumer).resume(Mockito.eq(emptySet())); } @Test @@ -3954,16 +3931,16 @@ public void shouldInitialiseNewStandbyTasks() { @Test public void shouldHandleRebalanceEvents() { final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.pause(assignment); - expectLastCall(); + taskManager.setMainConsumer(mockitoConsumer); + when(mockitoConsumer.assignment()).thenReturn(assignment); expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()); - replay(consumer, stateDirectory); + replay(stateDirectory); assertThat(taskManager.rebalanceInProgress(), is(false)); taskManager.handleRebalanceStart(emptySet()); assertThat(taskManager.rebalanceInProgress(), is(true)); taskManager.handleRebalanceComplete(); assertThat(taskManager.rebalanceInProgress(), is(false)); + Mockito.verify(mockitoConsumer).pause(assignment); } @Test @@ -4111,15 +4088,12 @@ public void shouldCommitViaConsumerIfEosDisabled() { final Map offsets = singletonMap(t1p1, new OffsetAndMetadata(0L, null)); task01.setCommittableOffsetsAndMetadata(offsets); task01.setCommitNeeded(); + taskManager.setMainConsumer(mockitoConsumer); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall(); - replay(consumer); - taskManager.commitAll(); - verify(consumer); + Mockito.verify(mockitoConsumer).commitSync(offsets); } @Test @@ -4159,6 +4133,7 @@ private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processing final Map offsetsT01, final Map offsetsT02) { final TaskManager taskManager = setUpTaskManager(processingMode, false); + taskManager.setMainConsumer(mockitoConsumer); final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); task01.setCommittableOffsetsAndMetadata(offsetsT01); @@ -4169,13 +4144,9 @@ private void shouldCommitViaProducerIfEosEnabled(final ProcessingMode processing task02.setCommitNeeded(); taskManager.addTask(task02); - reset(consumer); - expect(consumer.groupMetadata()).andStubReturn(new ConsumerGroupMetadata("appId")); - replay(consumer); + when(mockitoConsumer.groupMetadata()).thenReturn(new ConsumerGroupMetadata("appId")); taskManager.commitAll(); - - verify(consumer); } @Test @@ -4631,15 +4602,14 @@ public Set changelogPartitions() { } }; + taskManager.setMainConsumer(mockitoConsumer); + when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); 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); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -4870,11 +4840,10 @@ public void shouldThrowTaskMigratedExceptionOnCommitFailed() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task01.setCommittableOffsetsAndMetadata(offsets); task01.setCommitNeeded(); + taskManager.setMainConsumer(mockitoConsumer); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new CommitFailedException()); - replay(consumer); + doThrow(new CommitFailedException()).when(mockitoConsumer).commitSync(offsets); final TaskMigratedException thrown = assertThrows( TaskMigratedException.class, @@ -4899,11 +4868,9 @@ 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); + taskManager.setMainConsumer(mockitoConsumer); + + doThrow(new TimeoutException("KABOOM!")).doNothing().when(mockitoConsumer).commitSync(any(Map.class)); task00.setCommitNeeded(); @@ -4914,12 +4881,15 @@ public void shouldNotFailForTimeoutExceptionOnConsumerCommit() { assertThat(taskManager.commit(mkSet(task00, task01)), equalTo(1)); assertNull(task00.timeout); assertNull(task01.timeout); + + Mockito.verify(mockitoConsumer, times(2)).commitSync(any(Map.class)); } @Test public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() { final Tasks tasks = mock(Tasks.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_ALPHA, tasks, false); + taskManager.setMainConsumer(mockitoConsumer); final StreamsProducer producer = mock(StreamsProducer.class); when(activeTaskCreator.streamsProducerForTask(any(TaskId.class))).thenReturn(producer); @@ -4942,9 +4912,6 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() { task01.setCommittableOffsetsAndMetadata(offsetsT01); final StateMachineTask task02 = new StateMachineTask(taskId02, taskId02Partitions, true, stateManager); when(tasks.allTasks()).thenReturn(mkSet(task00, task01, task02)); - - expect(consumer.groupMetadata()).andStubReturn(null); - replay(consumer); task00.setCommitNeeded(); task01.setCommitNeeded(); @@ -4957,11 +4924,14 @@ public void shouldNotFailForTimeoutExceptionOnCommitWithEosAlpha() { exception.corruptedTasks(), equalTo(Collections.singleton(taskId00)) ); + + Mockito.verify(mockitoConsumer, times(2)).groupMetadata(); } @Test public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV2() { final TaskManager taskManager = setUpTaskManager(ProcessingMode.EXACTLY_ONCE_V2, false); + taskManager.setMainConsumer(mockitoConsumer); final StreamsProducer producer = mock(StreamsProducer.class); when(activeTaskCreator.threadProducer()).thenReturn(producer); @@ -4979,9 +4949,6 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV task01.setCommittableOffsetsAndMetadata(offsetsT01); final StateMachineTask task02 = new StateMachineTask(taskId02, taskId02Partitions, true, stateManager); - expect(consumer.groupMetadata()).andStubReturn(null); - replay(consumer); - task00.setCommitNeeded(); task01.setCommitNeeded(); @@ -4993,6 +4960,8 @@ public void shouldThrowTaskCorruptedExceptionForTimeoutExceptionOnCommitWithEosV exception.corruptedTasks(), equalTo(mkSet(taskId00, taskId01)) ); + + Mockito.verify(mockitoConsumer).groupMetadata(); } @Test @@ -5001,11 +4970,10 @@ public void shouldStreamsExceptionOnCommitError() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task01.setCommittableOffsetsAndMetadata(offsets); task01.setCommitNeeded(); + taskManager.setMainConsumer(mockitoConsumer); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new KafkaException()); - replay(consumer); + doThrow(new KafkaException()).when(mockitoConsumer).commitSync(offsets); final StreamsException thrown = assertThrows( StreamsException.class, @@ -5023,11 +4991,10 @@ public void shouldFailOnCommitFatal() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task01.setCommittableOffsetsAndMetadata(offsets); task01.setCommitNeeded(); + taskManager.setMainConsumer(mockitoConsumer); taskManager.addTask(task01); - consumer.commitSync(offsets); - expectLastCall().andThrow(new RuntimeException("KABOOM")); - replay(consumer); + doThrow(new RuntimeException("KABOOM")).when(mockitoConsumer).commitSync(offsets); final RuntimeException thrown = assertThrows( RuntimeException.class, @@ -5052,7 +5019,8 @@ public void suspend() { final Map> assignment = new HashMap<>(taskId00Assignment); assignment.putAll(taskId01Assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))).thenReturn(asList(task00, task01)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, Collections.emptyMap()); @@ -5063,6 +5031,7 @@ public void suspend() { assertThat(thrown.getCause().getMessage(), is("KABOOM!")); assertThat(task00.state(), is(Task.State.SUSPENDED)); assertThat(task01.state(), is(Task.State.SUSPENDED)); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -5078,7 +5047,7 @@ public void shouldConvertActiveTaskToStandbyTask() { when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(activeTask)); when(standbyTaskCreator.createStandbyTaskFromActive(Mockito.any(), Mockito.eq(taskId00Partitions))).thenReturn(standbyTask); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap()); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); @@ -5086,6 +5055,7 @@ public void shouldConvertActiveTaskToStandbyTask() { Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap())); Mockito.verify(standbyTaskCreator, times(2)).createTasks(Collections.emptyMap()); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test @@ -5102,13 +5072,14 @@ public void shouldConvertStandbyTaskToActiveTask() { when(activeTaskCreator.createActiveTaskFromStandby(Mockito.eq(standbyTask), Mockito.eq(taskId00Partitions), any())) .thenReturn(activeTask); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); taskManager.handleAssignment(taskId00Assignment, Collections.emptyMap()); Mockito.verify(activeTaskCreator, times(2)).createTasks(any(), Mockito.eq(emptyMap())); Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap()); + Mockito.verifyNoInteractions(mockitoConsumer); } @Test