Skip to content

KAFKA-14133: Migrate StateDirectory mock in TaskManagerTest to Mockito #13897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
TaklaGerguis wants to merge 2 commits into from
Closed

KAFKA-14133: Migrate StateDirectory mock in TaskManagerTest to Mockito #13897

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1c38b0a
KAFKA-14133: Migrate StateDirectory mock in TaskManagerTest to Mockito
TaklaGerguis Jun 21, 2023
1ff7536
Address comments
TaklaGerguis Jul 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 20 additions & 43 deletions streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,7 @@ public class TaskManagerTest {

@org.mockito.Mock
private InternalTopologyBuilder topologyBuilder;
@Mock(type = MockType.DEFAULT)
@org.mockito.Mock
private StateDirectory stateDirectory;
@org.mockito.Mock
private ChangelogReader changeLogReader;
Expand Down Expand Up @@ -1489,13 +1489,9 @@ public void shouldAddSubscribedTopicsFromAssignmentToTopologyMetadata() {

@Test
public void shouldNotLockAnythingIfStateDirIsEmpty() {
expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>()).once();

replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

verify(stateDirectory);
assertTrue(taskManager.lockedTaskDirectories().isEmpty());
Mockito.verify(stateDirectory).listNonEmptyTaskDirectories();
}

@Test
Expand All @@ -1508,10 +1504,7 @@ public void shouldTryToLockValidTaskDirsAtRebalanceStart() throws Exception {
taskId10.toString(),
"dummy"
);
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

verify(stateDirectory);
assertThat(taskManager.lockedTaskDirectories(), is(singleton(taskId01)));
}

Expand Down Expand Up @@ -1548,14 +1541,12 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
@Test
public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception {
expectLockObtainedFor(taskId00, taskId01, taskId02);
expectUnlockFor(taskId02);

makeTaskFolders(
taskId00.toString(), // active task
taskId01.toString(), // standby task
taskId02.toString() // unassigned but able to lock
);
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
Expand All @@ -1567,7 +1558,7 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception

taskManager.handleRebalanceComplete();
assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01)));
verify(stateDirectory);

@cadonna cadonna Jul 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please see my comment on expectUnlockFor(). You should replace this with expectUnlockFor(taskId02).

expectUnlockFor(taskId02);
}

@Test
Expand Down Expand Up @@ -1677,7 +1668,6 @@ private void computeOffsetSumAndVerify(final Map<TopicPartition, Long> changelog
final Map<TaskId, Long> expectedOffsetSums) throws Exception {
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
replay(stateDirectory);

taskManager.handleRebalanceStart(singleton("topic"));
final StateMachineTask restoringTask = handleAssignment(
Expand All @@ -1700,7 +1690,6 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception {

expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
replay(stateDirectory);

taskManager.handleRebalanceStart(singleton("topic"));
final StateMachineTask restoringTask = handleAssignment(
Expand All @@ -1724,8 +1713,6 @@ public void shouldComputeOffsetSumForUnassignedTaskWeCanLock() throws Exception
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
writeCheckpointFile(taskId00, changelogOffsets);

replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums));
Expand All @@ -1742,7 +1729,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForUninitializedTask() throw
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
writeCheckpointFile(taskId00, changelogOffsets);
replay(stateDirectory);

taskManager.handleRebalanceStart(singleton("topic"));
final StateMachineTask uninitializedTask = new StateMachineTask(taskId00, taskId00Partitions, true);
Expand All @@ -1766,7 +1752,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
writeCheckpointFile(taskId00, changelogOffsets);
replay(stateDirectory);

final StateMachineTask closedTask = new StateMachineTask(taskId00, taskId00Partitions, true);

Expand All @@ -1787,7 +1772,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep
public void shouldNotReportOffsetSumsForTaskWeCantLock() throws Exception {
expectLockFailedFor(taskId00);
makeTaskFolders(taskId00.toString());
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));
assertTrue(taskManager.lockedTaskDirectories().isEmpty());

Expand All @@ -1798,12 +1782,10 @@ public void shouldNotReportOffsetSumsForTaskWeCantLock() throws Exception {
public void shouldNotReportOffsetSumsAndReleaseLockForUnassignedTaskWithoutCheckpoint() throws Exception {
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
expect(stateDirectory.checkpointFileFor(taskId00)).andReturn(getCheckpointFile(taskId00));
replay(stateDirectory);
when(stateDirectory.checkpointFileFor(taskId00)).thenReturn(getCheckpointFile(taskId00));
taskManager.handleRebalanceStart(singleton("topic"));

assertTrue(taskManager.getTaskOffsetSums().isEmpty());
verify(stateDirectory);
}

@Test
Expand All @@ -1819,7 +1801,6 @@ public void shouldPinOffsetSumToLongMaxValueInCaseOfOverflow() throws Exception
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
writeCheckpointFile(taskId00, changelogOffsets);
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums));
Expand Down Expand Up @@ -1907,9 +1888,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
expectLockObtainedFor(taskId00, taskId01);

// The second attempt will return empty tasks.
makeTaskFolders();
expectLockObtainedFor();
replay(stateDirectory);

taskManager.handleRebalanceStart(emptySet());
assertThat(taskManager.lockedTaskDirectories(), Matchers.is(mkSet(taskId00, taskId01)));
Expand All @@ -1932,6 +1911,7 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
// The locked task map will not be cleared.
assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01)));

makeTaskFolders();
taskManager.handleRebalanceStart(emptySet());

assertThat(taskManager.lockedTaskDirectories(), is(emptySet()));
Expand Down Expand Up @@ -2161,7 +2141,7 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
}

@Test
public void shouldCleanAndReviveCorruptedStandbyTasksBeforeCommittingNonCorruptedTasks() {
public void shouldCleanAndReviveCorruptedStandbyTaskzBeforeCommittingNonCorruptedTaskz() {
final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class);
stateManager.markChangelogAsCorrupted(taskId00Partitions);
replay(stateManager);
Expand Down Expand Up @@ -2204,7 +2184,6 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
@Test
public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
final ProcessorStateManager stateManager = EasyMock.createNiceMock(ProcessorStateManager.class);
expect(stateDirectory.listNonEmptyTaskDirectories()).andStubReturn(new ArrayList<>());

final StateMachineTask corruptedActive = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager);

Expand All @@ -2224,7 +2203,7 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {

expect(consumer.assignment()).andStubReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));

replay(consumer, stateDirectory, stateManager);
replay(consumer, stateManager);

uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);

Expand Down Expand Up @@ -2848,7 +2827,7 @@ public void shouldCommitAllNeededTasksOnHandleRevocation() {
}

@Test
public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() {
public void shouldNotCommitOnHandleAssignmentIfNoTaskClozed() {
final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true);
final Map<TopicPartition, OffsetAndMetadata> offsets00 = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
task00.setCommittableOffsetsAndMetadata(offsets00);
Expand Down Expand Up @@ -2878,7 +2857,7 @@ public void shouldNotCommitOnHandleAssignmentIfNoTaskClosed() {
}

@Test
public void shouldNotCommitOnHandleAssignmentIfOnlyStandbyTaskClosed() {
public void shouldNotCommitOnHandleAzzignmentIfOnlyStandbyTaskClozed() {
final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true);
final Map<TopicPartition, OffsetAndMetadata> offsets00 = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
task00.setCommittableOffsetsAndMetadata(offsets00);
Expand Down Expand Up @@ -3428,8 +3407,7 @@ public void shouldHandleRebalanceEvents() {
expect(consumer.assignment()).andReturn(assignment);
consumer.pause(assignment);
expectLastCall();
expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>());
replay(consumer, stateDirectory);
replay(consumer);
assertThat(taskManager.rebalanceInProgress(), is(false));
taskManager.handleRebalanceStart(emptySet());
assertThat(taskManager.rebalanceInProgress(), is(true));
Expand Down Expand Up @@ -3543,15 +3521,15 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true);
final StateMachineTask task01 = new StateMachineTask(taskId01, taskId01Partitions, false);

makeTaskFolders(taskId00.toString(), task01.toString());
makeTaskFolders(taskId00.toString(), taskId01.toString());
expectLockObtainedFor(taskId00, taskId01);

@TaklaGerguis TaklaGerguis Jun 21, 2023

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was reported by Mockito as an unnecessary stubbing. I checked the code and it appears we never call this as part of the test. Let me know in case I am wrong.

@cadonna cadonna Jul 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is because there is a bug in the test.

Line

makeTaskFolders(taskId00.toString(), task01.toString());

should be

makeTaskFolders(taskId00.toString(), taskId01.toString());

Variable task01 holds a StateMachineTask and not a TaskId. Consequently, task01.toString() returns the string representation of a StateMachineTask and not of a TaskId which leads to a TaskIdFormatException in this call and skips the lock.

Please fix the bug and readd the stub.

expectRestoreToBeCompleted(consumer);
when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment)))
.thenReturn(singletonList(task00));
when(standbyTaskCreator.createTasks(taskId01Assignment))
.thenReturn(singletonList(task01));

replay(stateDirectory, consumer);
replay(consumer);

taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
Expand Down Expand Up @@ -4305,22 +4283,21 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
return allTasks;
}

private void expectLockObtainedFor(final TaskId... tasks) throws Exception {
private void expectLockObtainedFor(final TaskId... tasks) {
for (final TaskId task : tasks) {
expect(stateDirectory.lock(task)).andReturn(true).once();
when(stateDirectory.lock(task)).thenReturn(true);
}
}

private void expectLockFailedFor(final TaskId... tasks) throws Exception {
private void expectLockFailedFor(final TaskId... tasks) {
for (final TaskId task : tasks) {
expect(stateDirectory.lock(task)).andReturn(false).once();
when(stateDirectory.lock(task)).thenReturn(false);
}
}

private void expectUnlockFor(final TaskId... tasks) throws Exception {
private void expectUnlockFor(final TaskId... tasks) {
for (final TaskId task : tasks) {
stateDirectory.unlock(task);
expectLastCall();
Mockito.verify(stateDirectory, Mockito.atLeastOnce()).unlock(task);
}
}

Expand Down Expand Up @@ -4609,14 +4586,14 @@ private void makeTaskFolders(final String... names) throws Exception {
for (int i = 0; i < names.length; ++i) {
taskFolders.add(new TaskDirectory(testFolder.newFolder(names[i]), null));
}
expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(taskFolders).once();
when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(taskFolders);
}

private void writeCheckpointFile(final TaskId task, final Map<TopicPartition, Long> offsets) throws Exception {
final File checkpointFile = getCheckpointFile(task);
Files.createFile(checkpointFile.toPath());
new OffsetCheckpoint(checkpointFile).write(offsets);
expect(stateDirectory.checkpointFileFor(task)).andReturn(checkpointFile);
when(stateDirectory.checkpointFileFor(task)).thenReturn(checkpointFile);
}

private File getCheckpointFile(final TaskId task) {
Expand Down