From 3b151c795d57a0401b87d9388fca17e4f838dcd8 Mon Sep 17 00:00:00 2001 From: Christo Date: Thu, 25 Jan 2024 13:42:32 +0000 Subject: [PATCH 1/2] KAFKA-14133: Move consumer mock in TaskManagerTest to Mockito - part 2 --- .../processor/internals/TaskManagerTest.java | 424 +++++++++++------- 1 file changed, 256 insertions(+), 168 deletions(-) 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 ba1c91e7f7197..a742a16ef5b8a 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 @@ -58,7 +58,6 @@ import java.time.Duration; import java.util.ArrayList; -import org.easymock.EasyMock; import org.easymock.EasyMockRunner; import org.easymock.Mock; import org.easymock.MockType; @@ -112,7 +111,6 @@ 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.verify; import static org.hamcrest.CoreMatchers.hasItem; import static org.hamcrest.MatcherAssert.assertThat; @@ -133,6 +131,7 @@ import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.when; @@ -2016,13 +2015,14 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02))); handleAssignment(taskId00Assignment, taskId01Assignment, emptyMap()); - reset(consumer); - expectConsumerAssignmentPaused(consumer); - replay(consumer); taskManager.handleRebalanceComplete(); assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01))); verify(stateDirectory); + + final Set assignment = singleton(new TopicPartition("assignment", 0)); + Mockito.verify(mockitoConsumer).pause(assignment); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2210,6 +2210,9 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception { restoringTask.setChangelogOffsets(changelogOffsets); assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums)); + + final Set assignment = singleton(new TopicPartition("assignment", 0)); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2332,19 +2335,12 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() { task00.setCommittableOffsetsAndMetadata(offsets); // first `handleAssignment` - expectRestoreToBeCompleted(consumer); - when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - expectLastCall(); - - // `handleRevocation` - consumer.commitSync(offsets); - expectLastCall(); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); - // second `handleAssignment` - consumer.commitSync(offsets); - expectLastCall(); + when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2358,6 +2354,7 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() { assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2369,12 +2366,9 @@ public void closeClean() { } }; - // first `handleAssignment` - expectRestoreToBeCompleted(consumer); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); taskManager.handleRevocation(taskId00Partitions); @@ -2399,7 +2393,8 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); // `handleAssignment` - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); @@ -2412,11 +2407,11 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { expectLockObtainedFor(); replay(stateDirectory); + taskManager.setMainConsumer(mockitoConsumer); + taskManager.handleRebalanceStart(emptySet()); assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01))); - replay(consumer); - taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); @@ -2437,6 +2432,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { assertThat(taskManager.lockedTaskDirectories(), is(emptySet())); Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2455,15 +2451,14 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { task00.setCommittableOffsetsAndMetadata(offsets); // `handleAssignment` - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); // `handleAssignment` - consumer.commitSync(offsets); - expectLastCall(); doThrow(new RuntimeException("KABOOM!")).when(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2482,6 +2477,7 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { ); assertThat(thrown.getCause(), instanceOf(RuntimeException.class)); assertThat(thrown.getCause().getMessage(), is("KABOOM!")); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2529,10 +2525,13 @@ public void postCommit(final boolean enforceCheckpoint) { }; // `handleAssignment` - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(assignment) + .thenReturn(taskId00Partitions); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2548,7 +2547,7 @@ public void postCommit(final boolean enforceCheckpoint) { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2564,10 +2563,13 @@ public void suspend() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(assignment) + .thenReturn(taskId00Partitions); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2581,7 +2583,7 @@ public void suspend() { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2598,12 +2600,13 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { // `handleAssignment` when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(asList(corruptedTask, nonCorruptedTask)); - expectRestoreToBeCompleted(consumer); - expect(consumer.assignment()).andReturn(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(consumer); + + final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(consumerAssignment) + .thenReturn(taskId00Partitions); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); @@ -2618,7 +2621,9 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { assertThat(nonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet())); assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions)); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(consumerAssignment); + // check that we should not commit empty map either + Mockito.verify(mockitoConsumer, never()).commitSync(emptyMap()); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2637,8 +2642,9 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() { // `handleAssignment` when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(asList(corruptedTask, nonRunningNonCorruptedTask)); - expect(consumer.assignment()).andReturn(taskId00Partitions); - replay(consumer); + when(mockitoConsumer.assignment()).thenReturn(taskId00Partitions); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, emptyMap()); @@ -2650,7 +2656,6 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() { assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions)); assertFalse(nonRunningNonCorruptedTask.commitPrepared); - verify(consumer); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2732,9 +2737,10 @@ public Map prepareCommit() { .thenReturn(singleton(runningNonCorruptedActive)); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singleton(corruptedStandby)); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId01Assignment, taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2751,7 +2757,7 @@ public Map prepareCommit() { assertThat(corruptedStandby.commitPrepared, is(true)); assertThat(corruptedStandby.state(), is(Task.State.CREATED)); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2774,14 +2780,17 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(asList(corruptedActive, uncorruptedActive)); - expectRestoreToBeCompleted(consumer); - - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); + final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(consumerAssignment) + .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - replay(consumer, stateDirectory); + replay(stateDirectory); uncorruptedActive.setCommittableOffsetsAndMetadata(offsets); + taskManager.setMainConsumer(mockitoConsumer); + taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2800,7 +2809,7 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { assertThat(uncorruptedActive.commitCompleted, is(false)); assertThat(uncorruptedActive.state(), is(State.RUNNING)); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(consumerAssignment); } @Test @@ -2824,14 +2833,14 @@ public void markChangelogAsCorrupted(final Collection partitions when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(asList(corruptedActive, uncorruptedActive)); - expectRestoreToBeCompleted(consumer); + final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(consumerAssignment) + .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - consumer.commitSync(offsets); - expectLastCall().andThrow(new TimeoutException()); + doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(offsets); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2861,7 +2870,7 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(corruptedActive.state(), is(Task.State.CREATED)); assertThat(uncorruptedActive.state(), is(Task.State.CREATED)); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(consumerAssignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2899,16 +2908,17 @@ public void markChangelogAsCorrupted(final Collection partitions when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask)); - expectRestoreToBeCompleted(consumer); + final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(consumerAssignment) + .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); + when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata); doThrow(new TimeoutException()).when(producer).commitTransaction(offsets, groupMetadata); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); - - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -2944,7 +2954,7 @@ 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); + Mockito.verify(mockitoConsumer).resume(consumerAssignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions); } @@ -2978,16 +2988,17 @@ public void markChangelogAsCorrupted(final Collection partitions mkEntry(taskId02, taskId02Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(assignment) + .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(revokedActiveTask, unrevokedActiveTaskWithCommitNeeded, unrevokedActiveTaskWithoutCommitNeeded)); - expectLastCall(); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall().andThrow(new TimeoutException()); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - replay(consumer); + doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(expectedCommittedOffsets); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3000,6 +3011,8 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(revokedActiveTask.state(), is(State.SUSPENDED)); assertThat(unrevokedActiveTaskWithCommitNeeded.state(), is(State.CREATED)); assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3038,19 +3051,20 @@ public void markChangelogAsCorrupted(final Collection partitions mkEntry(taskId02, taskId02Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()) + .thenReturn(assignment) + .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(revokedActiveTask, unrevokedActiveTask, unrevokedActiveTaskWithoutCommitNeeded)); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); + when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata); doThrow(new TimeoutException()).when(producer).commitTransaction(expectedCommittedOffsets, groupMetadata); - expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); - - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3071,17 +3085,18 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00)); - consumer.commitSync(Collections.emptyMap()); - expectLastCall(); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(emptyMap(), taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3091,6 +3106,7 @@ public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { assertThat(task00.state(), is(Task.State.CLOSED)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3098,12 +3114,12 @@ public void shouldAddNonResumedSuspendedTasks() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - expectRestoreToBeCompleted(consumer); - // expect these calls twice (because we're going to tryToCompleteRestoration twice) - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3115,18 +3131,21 @@ public void shouldAddNonResumedSuspendedTasks() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThat(task01.state(), is(Task.State.RUNNING)); + // expect these calls twice (because we're going to tryToCompleteRestoration twice) Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap())); + Mockito.verify(mockitoConsumer, times(2)).assignment(); + Mockito.verify(mockitoConsumer, times(2)).resume(assignment); } @Test public void shouldUpdateInputPartitionsAfterRebalance() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - expectRestoreToBeCompleted(consumer); - // expect these calls twice (because we're going to tryToCompleteRestoration twice) - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3138,7 +3157,9 @@ public void shouldUpdateInputPartitionsAfterRebalance() { assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); assertEquals(newPartitionsSet, task00.inputPartitions()); - verify(consumer); + // expect these calls twice (because we're going to tryToCompleteRestoration twice) + Mockito.verify(mockitoConsumer, times(2)).resume(assignment); + Mockito.verify(mockitoConsumer, times(2)).assignment(); Mockito.verify(activeTaskCreator).createTasks(any(), Mockito.eq(emptyMap())); } @@ -3244,12 +3265,11 @@ public void shouldSuspendActiveTasksDuringRevocation() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task00.setCommittableOffsetsAndMetadata(offsets); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - consumer.commitSync(offsets); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3257,6 +3277,8 @@ public void shouldSuspendActiveTasksDuringRevocation() { taskManager.handleRevocation(taskId00Partitions); assertThat(task00.state(), is(Task.State.SUSPENDED)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3293,7 +3315,8 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo final Map> assignmentStandby = mkMap( mkEntry(taskId10, taskId10Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(task00, task01, task02)); @@ -3303,20 +3326,14 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo .thenReturn(singletonList(task10)); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); - expect(consumer.groupMetadata()).andReturn(groupMetadata); - producer.commitTransaction(expectedCommittedOffsets, groupMetadata); - expectLastCall(); + when(mockitoConsumer.groupMetadata()).thenReturn(groupMetadata); task00.committedOffsets(); - EasyMock.expectLastCall(); task01.committedOffsets(); - EasyMock.expectLastCall(); task02.committedOffsets(); - EasyMock.expectLastCall(); task10.committedOffsets(); - EasyMock.expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3331,6 +3348,9 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo assertThat(task01.commitNeeded, is(false)); assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); + Mockito.verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata); } @Test @@ -3364,16 +3384,15 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { final Map> assignmentStandby = mkMap( mkEntry(taskId10, taskId10Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(task00, task01, task02)); when(standbyTaskCreator.createTasks(assignmentStandby)) .thenReturn(singletonList(task10)); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3390,6 +3409,9 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { assertThat(task01.commitPrepared, is(true)); assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); + Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets); } @Test @@ -3404,12 +3426,13 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { final Map> assignmentActive = singletonMap(taskId00, taskId00Partitions); final Map> assignmentStandby = singletonMap(taskId10, taskId10Partitions); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00)); when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3420,6 +3443,8 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { assertThat(task00.commitNeeded, is(true)); assertThat(task10.commitPrepared, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3434,12 +3459,13 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { final Map> assignmentActive = singletonMap(taskId00, taskId00Partitions); final Map> assignmentStandby = singletonMap(taskId10, taskId10Partitions); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00)); when(standbyTaskCreator.createTasks(assignmentStandby)).thenReturn(singletonList(task10)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3449,6 +3475,8 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { taskManager.handleAssignment(assignmentActive, Collections.emptyMap()); assertThat(task00.commitNeeded, is(true)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3478,9 +3506,10 @@ public void suspend() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); @@ -3488,7 +3517,7 @@ public void suspend() { assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions)); assertThat(task00.state(), is(Task.State.SUSPENDED)); - verify(consumer); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3938,10 +3967,13 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() { @Test public void shouldInitializeNewActiveTasks() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); + when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3950,17 +3982,18 @@ public void shouldInitializeNewActiveTasks() { 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); + Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldInitialiseNewStandbyTasks() { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(emptyMap(), taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -3968,6 +4001,8 @@ public void shouldInitialiseNewStandbyTasks() { assertThat(task01.state(), is(Task.State.RUNNING)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.equalTo(singletonMap(taskId01, task01))); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3992,15 +4027,14 @@ public void shouldCommitActiveAndStandbyTasks() { task00.setCommittableOffsetsAndMetadata(offsets); final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)) .thenReturn(singletonList(task01)); - consumer.commitSync(offsets); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4014,6 +4048,9 @@ public void shouldCommitActiveAndStandbyTasks() { assertThat(taskManager.commitAll(), equalTo(2)); assertThat(task00.commitNeeded, is(false)); assertThat(task01.commitNeeded, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); + Mockito.verify(mockitoConsumer).commitSync(offsets); } @Test @@ -4036,15 +4073,14 @@ public void shouldCommitProvidedTasksIfNeeded() { mkEntry(taskId05, taskId05Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(Arrays.asList(task00, task01, task02)); when(standbyTaskCreator.createTasks(assignmentStandby)) .thenReturn(Arrays.asList(task03, task04, task05)); - consumer.commitSync(eq(emptyMap())); - - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4064,16 +4100,19 @@ public void shouldCommitProvidedTasksIfNeeded() { assertThat(task03.commitNeeded, is(false)); assertThat(task04.commitNeeded, is(true)); assertThat(task05.commitNeeded, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(Collections.emptyMap(), taskId00Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4084,6 +4123,8 @@ public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { assertThat(taskManager.commitAll(), equalTo(1)); assertThat(task00.commitNeeded, is(false)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4094,13 +4135,16 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw makeTaskFolders(taskId00.toString(), taskId01.toString()); expectDirectoryNotEmpty(taskId00, taskId01); expectLockObtainedFor(taskId00, taskId01); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)) .thenReturn(singletonList(task01)); - replay(stateDirectory, consumer); + replay(stateDirectory); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4122,6 +4166,8 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(-1) // sentinel indicating that nothing was done because a rebalance is in progress ); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4200,9 +4246,11 @@ public Map prepareCommit() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4214,6 +4262,8 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4225,10 +4275,11 @@ public Map prepareCommit() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(emptyMap(), taskId01Assignment); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4240,6 +4291,8 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4259,10 +4312,11 @@ public Map purgeableOffsets() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4278,6 +4332,8 @@ public Map purgeableOffsets() { inOrder.verify(adminClient).deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L))); inOrder.verify(adminClient).deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(17L))); inOrder.verifyNoMoreInteractions(); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4294,9 +4350,11 @@ public Map purgeableOffsets() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4311,20 +4369,23 @@ public Map purgeableOffsets() { // so it would fail verification if we invoke the admin client again. purgableOffsets.put(t1p1, 17L); taskManager.maybePurgeCommittedRecords(); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldIgnorePurgeDataErrors() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); final KafkaFutureImpl futureDeletedRecords = new KafkaFutureImpl<>(); final DeleteRecordsResult deleteRecordsResult = new DeleteRecordsResult(singletonMap(t1p1, futureDeletedRecords)); futureDeletedRecords.completeExceptionally(new Exception("KABOOM!")); when(adminClient.deleteRecords(any())).thenReturn(deleteRecordsResult); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.addTask(task00); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4335,6 +4396,8 @@ public void shouldIgnorePurgeDataErrors() { taskManager.maybePurgeCommittedRecords(); taskManager.maybePurgeCommittedRecords(); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4366,15 +4429,14 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { mkEntry(taskId10, taskId10Partitions) ); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(task00, task01, task02, task03)); when(standbyTaskCreator.createTasks(assignmentStandby)) .thenReturn(singletonList(task04)); - consumer.commitSync(expectedCommittedOffsets); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignmentActive, assignmentStandby); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4399,6 +4461,9 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { task04.setCommitRequested(); assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3)); + + Mockito.verify(mockitoConsumer).resume(assignment); + Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets); } @Test @@ -4410,10 +4475,12 @@ public void shouldProcessActiveTasks() { assignment.put(taskId00, taskId00Partitions); assignment.put(taskId01, taskId01Partitions); - expectRestoreToBeCompleted(consumer); + final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(consumerAssignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) .thenReturn(Arrays.asList(task00, task01)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4449,6 +4516,8 @@ public void shouldProcessActiveTasks() { // check that if there's no records processable, we would stop early assertThat(taskManager.process(3, time), is(5)); assertThat(taskManager.process(3, time), is(0)); + + Mockito.verify(mockitoConsumer).resume(consumerAssignment); } @Test @@ -4523,9 +4592,11 @@ public boolean process(final long wallClockTime) { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4536,6 +4607,8 @@ public boolean process(final long wallClockTime) { task00.addRecords(partition, singletonList(getConsumerRecord(partition, 0L))); assertThrows(TaskMigratedException.class, () -> taskManager.process(1, time)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4547,10 +4620,12 @@ public boolean process(final long wallClockTime) { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4564,6 +4639,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")); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4575,9 +4652,11 @@ public boolean maybePunctuateStreamTime() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4585,6 +4664,8 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(TaskMigratedException.class, () -> taskManager.punctuate()); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4596,9 +4677,11 @@ public boolean maybePunctuateStreamTime() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4606,6 +4689,8 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(KafkaException.class, () -> taskManager.punctuate()); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4622,9 +4707,11 @@ public boolean maybePunctuateSystemTime() { } }; - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - replay(consumer); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(taskId00Assignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); @@ -4633,6 +4720,8 @@ public boolean maybePunctuateSystemTime() { // one for stream and one for system time assertThat(taskManager.punctuate(), equalTo(2)); + + Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4660,12 +4749,11 @@ public void shouldHaveRemainingPartitionsUncleared() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task00.setCommittableOffsetsAndMetadata(offsets); - expectRestoreToBeCompleted(consumer); + final Set assignment = singleton(new TopicPartition("assignment", 0)); + when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); - consumer.commitSync(offsets); - expectLastCall(); - replay(consumer); + taskManager.setMainConsumer(mockitoConsumer); try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) { appender.setClassLoggerToDebug(TaskManager.class); @@ -4684,6 +4772,8 @@ public void shouldHaveRemainingPartitionsUncleared() { "condition of consumer detecting the heartbeat failure, or the " + "tasks have been cleaned up by the handleAssignment callback.") ); + + Mockito.verify(mockitoConsumer).resume(assignment); } } @@ -4821,8 +4911,10 @@ private Map handleAssignment(final Map assignment = singleton(new TopicPartition("assignment", 0)); + lenient().when(mockitoConsumer.assignment()).thenReturn(assignment); + + taskManager.setMainConsumer(mockitoConsumer); taskManager.handleAssignment(allActiveTasksAssignment, standbyAssignment); taskManager.tryToCompleteRestoration(time.milliseconds(), null); @@ -5133,13 +5225,9 @@ public void shouldListNotPausedTasks() { topologyMetadata.pauseTopology(UNNAMED_TOPOLOGY); assertEquals(taskManager.notPausedTasks().size(), 0); - } - private static void expectRestoreToBeCompleted(final Consumer consumer) { final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.resume(assignment); - expectLastCall(); + Mockito.verify(mockitoConsumer).resume(assignment); } private static KafkaFutureImpl completedFuture() { From 2c621e0d61474c84295c0d30d01c249f6db74cf1 Mon Sep 17 00:00:00 2001 From: Christo Date: Wed, 6 Mar 2024 17:14:42 +0000 Subject: [PATCH 2/2] Address comments from review --- .../processor/internals/TaskManagerTest.java | 173 +++--------------- 1 file changed, 30 insertions(+), 143 deletions(-) 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 a742a16ef5b8a..681e69d300487 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 @@ -184,6 +184,7 @@ public class TaskManagerTest { private final TaskId taskId10 = new TaskId(1, 0); private final TopicPartition t2p0 = new TopicPartition(topic2, 0); private final Set taskId10Partitions = mkSet(t2p0); + private final Set assignment = singleton(new TopicPartition("assignment", 0)); final java.util.function.Consumer> noOpResetter = partitions -> { }; @@ -2020,9 +2021,7 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01))); verify(stateDirectory); - final Set assignment = singleton(new TopicPartition("assignment", 0)); Mockito.verify(mockitoConsumer).pause(assignment); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2210,9 +2209,6 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception { restoringTask.setChangelogOffsets(changelogOffsets); assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums)); - - final Set assignment = singleton(new TopicPartition("assignment", 0)); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2335,7 +2331,6 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() { task00.setCommittableOffsetsAndMetadata(offsets); // first `handleAssignment` - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -2354,7 +2349,6 @@ public void shouldCloseActiveUnassignedSuspendedTasksWhenClosingRevokedTasks() { assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2393,7 +2387,6 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); // `handleAssignment` - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); @@ -2432,7 +2425,6 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception { assertThat(taskManager.lockedTaskDirectories(), is(emptySet())); Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2451,7 +2443,6 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { task00.setCommittableOffsetsAndMetadata(offsets); // `handleAssignment` - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -2477,7 +2468,6 @@ public void shouldThrowWhenHandlingClosingTasksOnProducerCloseError() { ); assertThat(thrown.getCause(), instanceOf(RuntimeException.class)); assertThat(thrown.getCause().getMessage(), is("KABOOM!")); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -2525,7 +2515,6 @@ public void postCommit(final boolean enforceCheckpoint) { }; // `handleAssignment` - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) .thenReturn(assignment) .thenReturn(taskId00Partitions); @@ -2547,7 +2536,6 @@ public void postCommit(final boolean enforceCheckpoint) { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2563,7 +2551,6 @@ public void suspend() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) .thenReturn(assignment) .thenReturn(taskId00Partitions); @@ -2583,7 +2570,6 @@ public void suspend() { assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00))); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2594,21 +2580,20 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { final StateMachineTask corruptedTask = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final StateMachineTask nonCorruptedTask = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); - final Map> assignment = new HashMap<>(taskId00Assignment); - assignment.putAll(taskId01Assignment); + final Map> firstAssignment = new HashMap<>(taskId00Assignment); + firstAssignment.putAll(taskId01Assignment); // `handleAssignment` - when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) + when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment))) .thenReturn(asList(corruptedTask, nonCorruptedTask)); - final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) - .thenReturn(consumerAssignment) + .thenReturn(assignment) .thenReturn(taskId00Partitions); taskManager.setMainConsumer(mockitoConsumer); - taskManager.handleAssignment(assignment, emptyMap()); + taskManager.handleAssignment(firstAssignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true)); assertThat(nonCorruptedTask.state(), is(Task.State.RUNNING)); @@ -2621,7 +2606,6 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() { assertThat(nonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet())); assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions)); - Mockito.verify(mockitoConsumer).resume(consumerAssignment); // check that we should not commit empty map either Mockito.verify(mockitoConsumer, never()).commitSync(emptyMap()); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); @@ -2737,7 +2721,6 @@ public Map prepareCommit() { .thenReturn(singleton(runningNonCorruptedActive)); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singleton(corruptedStandby)); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); taskManager.setMainConsumer(mockitoConsumer); @@ -2757,7 +2740,6 @@ public Map prepareCommit() { assertThat(corruptedStandby.commitPrepared, is(true)); assertThat(corruptedStandby.state(), is(Task.State.CREATED)); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2774,15 +2756,14 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { uncorruptedActive.setCommitNeeded(); // handleAssignment - final Map> assignment = new HashMap<>(); - assignment.putAll(taskId00Assignment); - assignment.putAll(taskId01Assignment); - when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) + final Map> firstAssignement = new HashMap<>(); + firstAssignement.putAll(taskId00Assignment); + firstAssignement.putAll(taskId01Assignment); + when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignement))) .thenReturn(asList(corruptedActive, uncorruptedActive)); - final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) - .thenReturn(consumerAssignment) + .thenReturn(assignment) .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); replay(stateDirectory); @@ -2791,7 +2772,7 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { taskManager.setMainConsumer(mockitoConsumer); - taskManager.handleAssignment(assignment, emptyMap()); + taskManager.handleAssignment(firstAssignement, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(uncorruptedActive.state(), is(Task.State.RUNNING)); @@ -2809,7 +2790,6 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() { assertThat(uncorruptedActive.commitCompleted, is(false)); assertThat(uncorruptedActive.state(), is(State.RUNNING)); - Mockito.verify(mockitoConsumer).resume(consumerAssignment); } @Test @@ -2827,22 +2807,21 @@ public void markChangelogAsCorrupted(final Collection partitions uncorruptedActive.setCommittableOffsetsAndMetadata(offsets); // handleAssignment - final Map> assignment = new HashMap<>(); - assignment.putAll(taskId00Assignment); - assignment.putAll(taskId01Assignment); - when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) + final Map> firstAssignment = new HashMap<>(); + firstAssignment.putAll(taskId00Assignment); + firstAssignment.putAll(taskId01Assignment); + when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment))) .thenReturn(asList(corruptedActive, uncorruptedActive)); - final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) - .thenReturn(consumerAssignment) + .thenReturn(assignment) .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); doThrow(new TimeoutException()).when(mockitoConsumer).commitSync(offsets); taskManager.setMainConsumer(mockitoConsumer); - taskManager.handleAssignment(assignment, emptyMap()); + taskManager.handleAssignment(firstAssignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(uncorruptedActive.state(), is(Task.State.RUNNING)); @@ -2870,7 +2849,6 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(corruptedActive.state(), is(Task.State.CREATED)); assertThat(uncorruptedActive.state(), is(Task.State.CREATED)); - Mockito.verify(mockitoConsumer).resume(consumerAssignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00Partitions); } @@ -2902,15 +2880,14 @@ public void markChangelogAsCorrupted(final Collection partitions uncorruptedActiveTask.setCommittableOffsetsAndMetadata(offsets); // handleAssignment - final Map> assignment = new HashMap<>(); - assignment.putAll(taskId00Assignment); - assignment.putAll(taskId01Assignment); - when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) + final Map> firstAssignment = new HashMap<>(); + firstAssignment.putAll(taskId00Assignment); + firstAssignment.putAll(taskId01Assignment); + when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment))) .thenReturn(asList(corruptedActiveTask, uncorruptedActiveTask)); - final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) - .thenReturn(consumerAssignment) + .thenReturn(assignment) .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions)); final ConsumerGroupMetadata groupMetadata = new ConsumerGroupMetadata("appId"); @@ -2920,7 +2897,7 @@ public void markChangelogAsCorrupted(final Collection partitions taskManager.setMainConsumer(mockitoConsumer); - taskManager.handleAssignment(assignment, emptyMap()); + taskManager.handleAssignment(firstAssignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(uncorruptedActiveTask.state(), is(Task.State.RUNNING)); @@ -2954,7 +2931,6 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(uncorruptedActiveTask.state(), is(Task.State.CREATED)); assertThat(corruptedTaskChangelogMarkedAsCorrupted.get(), is(true)); assertThat(uncorruptedTaskChangelogMarkedAsCorrupted.get(), is(true)); - Mockito.verify(mockitoConsumer).resume(consumerAssignment); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions); } @@ -2988,7 +2964,6 @@ public void markChangelogAsCorrupted(final Collection partitions mkEntry(taskId02, taskId02Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) .thenReturn(assignment) .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); @@ -3011,8 +2986,6 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(revokedActiveTask.state(), is(State.SUSPENDED)); assertThat(unrevokedActiveTaskWithCommitNeeded.state(), is(State.CREATED)); assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3051,7 +3024,6 @@ public void markChangelogAsCorrupted(final Collection partitions mkEntry(taskId02, taskId02Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()) .thenReturn(assignment) .thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions)); @@ -3085,14 +3057,12 @@ public void markChangelogAsCorrupted(final Collection partitions assertThat(unrevokedActiveTaskWithoutCommitNeeded.state(), is(State.RUNNING)); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId00ChangelogPartitions); Mockito.verify(stateManager).markChangelogAsCorrupted(taskId01ChangelogPartitions); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00)); @@ -3106,7 +3076,6 @@ public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() { assertThat(task00.state(), is(Task.State.CLOSED)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap()); - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3114,7 +3083,6 @@ public void shouldAddNonResumedSuspendedTasks() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); @@ -3141,7 +3109,6 @@ public void shouldAddNonResumedSuspendedTasks() { public void shouldUpdateInputPartitionsAfterRebalance() { final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -3265,7 +3232,6 @@ public void shouldSuspendActiveTasksDuringRevocation() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task00.setCommittableOffsetsAndMetadata(offsets); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -3277,8 +3243,6 @@ public void shouldSuspendActiveTasksDuringRevocation() { taskManager.handleRevocation(taskId00Partitions); assertThat(task00.state(), is(Task.State.SUSPENDED)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3315,7 +3279,6 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo final Map> assignmentStandby = mkMap( mkEntry(taskId10, taskId10Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) @@ -3349,7 +3312,6 @@ public void shouldCommitAllActiveTasksThatNeedCommittingOnHandleRevocationWithEo assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(producer).commitTransaction(expectedCommittedOffsets, groupMetadata); } @@ -3384,7 +3346,6 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { final Map> assignmentStandby = mkMap( mkEntry(taskId10, taskId10Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) @@ -3410,7 +3371,6 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() { assertThat(task02.commitPrepared, is(false)); assertThat(task10.commitPrepared, is(false)); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets); } @@ -3426,7 +3386,6 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { final Map> assignmentActive = singletonMap(taskId00, taskId00Partitions); final Map> assignmentStandby = singletonMap(taskId10, taskId10Partitions); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00)); @@ -3443,8 +3402,6 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() { assertThat(task00.commitNeeded, is(true)); assertThat(task10.commitPrepared, is(false)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3459,7 +3416,6 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { final Map> assignmentActive = singletonMap(taskId00, taskId00Partitions); final Map> assignmentStandby = singletonMap(taskId10, taskId10Partitions); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))).thenReturn(singleton(task00)); @@ -3475,8 +3431,6 @@ public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() { taskManager.handleAssignment(assignmentActive, Collections.emptyMap()); assertThat(task00.commitNeeded, is(true)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3506,7 +3460,6 @@ public void suspend() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); taskManager.setMainConsumer(mockitoConsumer); @@ -3516,8 +3469,6 @@ public void suspend() { assertThrows(RuntimeException.class, () -> taskManager.handleRevocation(taskId00Partitions)); assertThat(task00.state(), is(Task.State.SUSPENDED)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -3967,7 +3918,6 @@ public void shouldShutDownStateUpdaterAndAddRemovedTasksToTaskRegistry() { @Test public void shouldInitializeNewActiveTasks() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) @@ -3989,7 +3939,6 @@ public void shouldInitializeNewActiveTasks() { public void shouldInitialiseNewStandbyTasks() { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); @@ -4001,13 +3950,10 @@ public void shouldInitialiseNewStandbyTasks() { assertThat(task01.state(), is(Task.State.RUNNING)); assertThat(taskManager.activeTaskMap(), Matchers.anEmptyMap()); assertThat(taskManager.standbyTaskMap(), Matchers.equalTo(singletonMap(taskId01, task01))); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldHandleRebalanceEvents() { - final Set assignment = singleton(new TopicPartition("assignment", 0)); taskManager.setMainConsumer(mockitoConsumer); when(mockitoConsumer.assignment()).thenReturn(assignment); expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()); @@ -4027,7 +3973,6 @@ public void shouldCommitActiveAndStandbyTasks() { task00.setCommittableOffsetsAndMetadata(offsets); final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); @@ -4049,7 +3994,6 @@ public void shouldCommitActiveAndStandbyTasks() { assertThat(task00.commitNeeded, is(false)); assertThat(task01.commitNeeded, is(false)); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(mockitoConsumer).commitSync(offsets); } @@ -4073,7 +4017,6 @@ public void shouldCommitProvidedTasksIfNeeded() { mkEntry(taskId05, taskId05Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(Arrays.asList(task00, task01, task02)); @@ -4100,15 +4043,12 @@ public void shouldCommitProvidedTasksIfNeeded() { assertThat(task03.commitNeeded, is(false)); assertThat(task04.commitNeeded, is(true)); assertThat(task05.commitNeeded, is(false)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId00Assignment)).thenReturn(singletonList(task00)); @@ -4123,8 +4063,6 @@ public void shouldNotCommitOffsetsIfOnlyStandbyTasksAssigned() { assertThat(taskManager.commitAll(), equalTo(1)); assertThat(task00.commitNeeded, is(false)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4135,7 +4073,6 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw makeTaskFolders(taskId00.toString(), taskId01.toString()); expectDirectoryNotEmpty(taskId00, taskId01); expectLockObtainedFor(taskId00, taskId01); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); @@ -4166,8 +4103,6 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(-1) // sentinel indicating that nothing was done because a rebalance is in progress ); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4246,7 +4181,6 @@ public Map prepareCommit() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4262,8 +4196,6 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4275,7 +4207,6 @@ public Map prepareCommit() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01)); @@ -4291,8 +4222,6 @@ public Map prepareCommit() { final RuntimeException thrown = assertThrows(RuntimeException.class, () -> taskManager.commitAll()); assertThat(thrown.getMessage(), equalTo("opsh.")); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4312,7 +4241,6 @@ public Map purgeableOffsets() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4332,8 +4260,6 @@ public Map purgeableOffsets() { inOrder.verify(adminClient).deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(5L))); inOrder.verify(adminClient).deleteRecords(singletonMap(t1p1, RecordsToDelete.beforeOffset(17L))); inOrder.verifyNoMoreInteractions(); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4350,7 +4276,6 @@ public Map purgeableOffsets() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4369,15 +4294,12 @@ public Map purgeableOffsets() { // so it would fail verification if we invoke the admin client again. purgableOffsets.put(t1p1, 17L); taskManager.maybePurgeCommittedRecords(); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test public void shouldIgnorePurgeDataErrors() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); final KafkaFutureImpl futureDeletedRecords = new KafkaFutureImpl<>(); @@ -4396,8 +4318,6 @@ public void shouldIgnorePurgeDataErrors() { taskManager.maybePurgeCommittedRecords(); taskManager.maybePurgeCommittedRecords(); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4429,7 +4349,6 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { mkEntry(taskId10, taskId10Partitions) ); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(assignmentActive))) .thenReturn(asList(task00, task01, task02, task03)); @@ -4462,7 +4381,6 @@ public void shouldMaybeCommitAllActiveTasksThatNeedCommit() { assertThat(taskManager.maybeCommitActiveTasksPerUserRequested(), equalTo(3)); - Mockito.verify(mockitoConsumer).resume(assignment); Mockito.verify(mockitoConsumer).commitSync(expectedCommittedOffsets); } @@ -4471,18 +4389,17 @@ public void shouldProcessActiveTasks() { final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager); final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); - final Map> assignment = new HashMap<>(); - assignment.put(taskId00, taskId00Partitions); - assignment.put(taskId01, taskId01Partitions); + final Map> firstAssignment = new HashMap<>(); + firstAssignment.put(taskId00, taskId00Partitions); + firstAssignment.put(taskId01, taskId01Partitions); - final Set consumerAssignment = singleton(new TopicPartition("assignment", 0)); - when(mockitoConsumer.assignment()).thenReturn(consumerAssignment); - when(activeTaskCreator.createTasks(any(), Mockito.eq(assignment))) + when(mockitoConsumer.assignment()).thenReturn(assignment); + when(activeTaskCreator.createTasks(any(), Mockito.eq(firstAssignment))) .thenReturn(Arrays.asList(task00, task01)); taskManager.setMainConsumer(mockitoConsumer); - taskManager.handleAssignment(assignment, emptyMap()); + taskManager.handleAssignment(firstAssignment, emptyMap()); assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true)); assertThat(task00.state(), is(Task.State.RUNNING)); @@ -4516,8 +4433,6 @@ public void shouldProcessActiveTasks() { // check that if there's no records processable, we would stop early assertThat(taskManager.process(3, time), is(5)); assertThat(taskManager.process(3, time), is(0)); - - Mockito.verify(mockitoConsumer).resume(consumerAssignment); } @Test @@ -4592,7 +4507,6 @@ public boolean process(final long wallClockTime) { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4607,8 +4521,6 @@ public boolean process(final long wallClockTime) { task00.addRecords(partition, singletonList(getConsumerRecord(partition, 0L))); assertThrows(TaskMigratedException.class, () -> taskManager.process(1, time)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4620,7 +4532,6 @@ public boolean process(final long wallClockTime) { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))) .thenReturn(singletonList(task00)); @@ -4639,8 +4550,6 @@ public boolean process(final long wallClockTime) { assertThat(exception.taskId().isPresent(), is(true)); assertThat(exception.taskId().get(), is(taskId00)); assertThat(exception.getCause().getMessage(), is("oops")); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4652,7 +4561,6 @@ public boolean maybePunctuateStreamTime() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4664,8 +4572,6 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(TaskMigratedException.class, () -> taskManager.punctuate()); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4677,7 +4583,6 @@ public boolean maybePunctuateStreamTime() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4689,8 +4594,6 @@ public boolean maybePunctuateStreamTime() { assertThat(task00.state(), is(Task.State.RUNNING)); assertThrows(KafkaException.class, () -> taskManager.punctuate()); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4707,7 +4610,6 @@ public boolean maybePunctuateSystemTime() { } }; - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4720,8 +4622,6 @@ public boolean maybePunctuateSystemTime() { // one for stream and one for system time assertThat(taskManager.punctuate(), equalTo(2)); - - Mockito.verify(mockitoConsumer).resume(assignment); } @Test @@ -4749,7 +4649,6 @@ public void shouldHaveRemainingPartitionsUncleared() { final Map offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null)); task00.setCommittableOffsetsAndMetadata(offsets); - final Set assignment = singleton(new TopicPartition("assignment", 0)); when(mockitoConsumer.assignment()).thenReturn(assignment); when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00)); @@ -4772,8 +4671,6 @@ public void shouldHaveRemainingPartitionsUncleared() { "condition of consumer detecting the heartbeat failure, or the " + "tasks have been cleaned up by the handleAssignment callback.") ); - - Mockito.verify(mockitoConsumer).resume(assignment); } } @@ -4911,7 +4808,6 @@ private Map handleAssignment(final Map assignment = singleton(new TopicPartition("assignment", 0)); lenient().when(mockitoConsumer.assignment()).thenReturn(assignment); taskManager.setMainConsumer(mockitoConsumer); @@ -4962,12 +4858,6 @@ private void expectDirectoryNotEmpty(final TaskId... tasks) { } } - private static void expectConsumerAssignmentPaused(final Consumer consumer) { - final Set assignment = singleton(new TopicPartition("assignment", 0)); - expect(consumer.assignment()).andReturn(assignment); - consumer.pause(assignment); - } - @Test public void shouldThrowTaskMigratedExceptionOnCommitFailed() { final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager); @@ -5225,9 +5115,6 @@ public void shouldListNotPausedTasks() { topologyMetadata.pauseTopology(UNNAMED_TOPOLOGY); assertEquals(taskManager.notPausedTasks().size(), 0); - - final Set assignment = singleton(new TopicPartition("assignment", 0)); - Mockito.verify(mockitoConsumer).resume(assignment); } private static KafkaFutureImpl completedFuture() {