Skip to content

KAFKA-14133: Move stateDirectory mock in TaskManagerTest to Mockito #15254

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

Merged
cadonna merged 4 commits into from
Mar 22, 2024
Merged

KAFKA-14133: Move stateDirectory mock in TaskManagerTest to Mockito #15254

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6c657fc
KAFKA-14133: Move stateDirectory mock in TaskManagerTest to Mockito
clolov Jan 24, 2024
fc370f3
Rebase and remove references to EasyMock
clolov Mar 18, 2024
fd43760
Addressing review comments
clolov Mar 18, 2024
9974336
Address comments from second round of review
clolov Mar 18, 2024
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
98 changes: 29 additions & 69 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 @@ -58,9 +58,6 @@
import java.time.Duration;
import java.util.ArrayList;

import org.easymock.EasyMockRunner;
import org.easymock.Mock;
import org.easymock.MockType;
import org.hamcrest.Matchers;
import org.junit.Before;
import org.junit.Rule;
Expand All @@ -71,8 +68,8 @@
import org.mockito.InOrder;
import org.mockito.Mockito;
import org.mockito.junit.MockitoJUnit;
import org.mockito.junit.MockitoJUnitRunner;
import org.mockito.junit.MockitoRule;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;

import java.io.File;
Expand Down Expand Up @@ -107,11 +104,6 @@
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask;
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask;
import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.easymock.EasyMock.replay;
import static org.easymock.EasyMock.verify;
import static org.hamcrest.CoreMatchers.hasItem;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.empty;
Expand All @@ -137,8 +129,7 @@
import static org.mockito.Mockito.when;
import static org.mockito.Mockito.mock;

@MockitoSettings(strictness = Strictness.STRICT_STUBS)
@RunWith(EasyMockRunner.class)
@RunWith(MockitoJUnitRunner.StrictStubs.class)
public class TaskManagerTest {

private final String topic1 = "topic1";
Expand Down Expand Up @@ -190,7 +181,7 @@ public class TaskManagerTest {

@org.mockito.Mock
private InternalTopologyBuilder topologyBuilder;
@Mock(type = MockType.DEFAULT)
@org.mockito.Mock

@ijuma ijuma Mar 18, 2024

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.

Can we now import this and avoid the fully qualified name?

@cadonna cadonna Mar 18, 2024

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.

@clolov If you resolve Ismaels comment, could you also please remove the Mockito prefix for other method calls like verify()?

@clolov clolov Mar 18, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Hopefully done on both accounts 😊 !

private StateDirectory stateDirectory;
@org.mockito.Mock
private ChangelogReader changeLogReader;
Expand Down Expand Up @@ -785,7 +776,6 @@ public void shouldNotReturnStateUpdaterTasksInOwnedTasks() {
final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);

when(stateUpdater.getTasks()).thenReturn(mkSet(standbyTask));

@clolov clolov Mar 18, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This was reported as unnecessary stubbing by Mockito

when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId03, activeTask)));
assertEquals(taskManager.allOwnedTasks(), mkMap(mkEntry(taskId03, activeTask)));
}
Expand Down Expand Up @@ -1920,12 +1910,10 @@ public void shouldAddSubscribedTopicsFromAssignmentToTopologyMetadata() {

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

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

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

Expand All @@ -1940,25 +1928,21 @@ public void shouldTryToLockValidTaskDirsAtRebalanceStart() throws Exception {
taskId10.toString(),
"dummy"
);
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

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

@Test
public void shouldUnlockEmptyDirsAtRebalanceStart() throws Exception {
expectLockObtainedFor(taskId01, taskId10);
expectDirectoryNotEmpty(taskId01);
expect(stateDirectory.directoryForTaskIsEmpty(taskId10)).andReturn(true);
expectUnlockFor(taskId10);
when(stateDirectory.directoryForTaskIsEmpty(taskId10)).thenReturn(true);

makeTaskFolders(taskId01.toString(), taskId10.toString());
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

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

Expand Down Expand Up @@ -1992,14 +1976,12 @@ public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception {
expectLockObtainedFor(taskId00, taskId01, taskId02);
expectDirectoryNotEmpty(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 @@ -2008,8 +1990,8 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception

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

Mockito.verify(stateDirectory).unlock(taskId02);
Mockito.verify(consumer).pause(assignment);
}

Expand All @@ -2034,14 +2016,12 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
when(tasks.allTasks()).thenReturn(mkSet(runningStatefulTask));
expectLockObtainedFor(taskId00, taskId01, taskId02, taskId03);
expectDirectoryNotEmpty(taskId00, taskId01, taskId02, taskId03);
expectUnlockFor(taskId03);
makeTaskFolders(
taskId00.toString(),
taskId01.toString(),
taskId02.toString(),
taskId03.toString()
);
replay(stateDirectory);

final Set<TopicPartition> assigned = mkSet(t1p0, t1p1, t1p2);
when(consumer.assignment()).thenReturn(assigned);
Expand All @@ -2050,7 +2030,7 @@ public void shouldReleaseLockForUnassignedTasksAfterRebalanceWithStateUpdater()
taskManager.handleRebalanceComplete();

Mockito.verify(consumer).pause(mkSet(t1p1, t1p2));
verify(stateDirectory);
Mockito.verify(stateDirectory).unlock(taskId03);
assertThat(taskManager.lockedTaskDirectories(), is(mkSet(taskId00, taskId01, taskId02)));
}

Expand Down Expand Up @@ -2089,7 +2069,6 @@ public void shouldComputeOffsetSumForRestoringActiveTaskWithStateUpdater() throw
final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
when(stateUpdater.getTasks()).thenReturn(mkSet(restoringStatefulTask));
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

assertThat(taskManager.getTaskOffsetSums(), is(mkMap(mkEntry(taskId00, changelogOffset))));
Expand All @@ -2108,7 +2087,6 @@ public void shouldComputeOffsetSumForRestoringStandbyTaskWithStateUpdater() thro
final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
when(stateUpdater.getTasks()).thenReturn(mkSet(restoringStandbyTask));
replay(stateDirectory);
taskManager.handleRebalanceStart(singleton("topic"));

assertThat(taskManager.getTaskOffsetSums(), is(mkMap(mkEntry(taskId00, changelogOffset))));
Expand Down Expand Up @@ -2162,7 +2140,6 @@ private void computeOffsetSumAndVerify(final Map<TopicPartition, Long> changelog
expectLockObtainedFor(taskId00);
expectDirectoryNotEmpty(taskId00);
makeTaskFolders(taskId00.toString());
replay(stateDirectory);

taskManager.handleRebalanceStart(singleton("topic"));
final StateMachineTask restoringTask = handleAssignment(
Expand All @@ -2186,7 +2163,6 @@ public void shouldComputeOffsetSumForStandbyTask() throws Exception {
expectLockObtainedFor(taskId00);
expectDirectoryNotEmpty(taskId00);
makeTaskFolders(taskId00.toString());
replay(stateDirectory);

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

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

assertThat(taskManager.getTaskOffsetSums(), is(expectedOffsetSums));
Expand All @@ -2228,7 +2203,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, stateManager);
Expand All @@ -2252,7 +2226,6 @@ public void shouldComputeOffsetSumFromCheckpointFileForClosedTask() throws Excep
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
writeCheckpointFile(taskId00, changelogOffsets);
replay(stateDirectory);

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

Expand All @@ -2273,7 +2246,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 @@ -2285,12 +2257,10 @@ public void shouldNotReportOffsetSumsAndReleaseLockForUnassignedTaskWithoutCheck
expectLockObtainedFor(taskId00);
makeTaskFolders(taskId00.toString());
expectDirectoryNotEmpty(taskId00);
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 @@ -2306,7 +2276,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 @@ -2375,15 +2344,18 @@ public void shouldCloseActiveTasksWhenHandlingLostTasks() throws Exception {
when(activeTaskCreator.createTasks(any(), Mockito.eq(taskId00Assignment))).thenReturn(singletonList(task00));
when(standbyTaskCreator.createTasks(taskId01Assignment)).thenReturn(singletonList(task01));

makeTaskFolders(taskId00.toString(), taskId01.toString());
final ArrayList<TaskDirectory> taskFolders = new ArrayList<>(2);
taskFolders.add(new TaskDirectory(testFolder.newFolder(taskId00.toString()), null));
taskFolders.add(new TaskDirectory(testFolder.newFolder(taskId01.toString()), null));

when(stateDirectory.listNonEmptyTaskDirectories())
.thenReturn(taskFolders)
// The second attempt will return empty tasks.

@cadonna cadonna Mar 4, 2024

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.

nit: This comment is not needed. Please remove.

.thenReturn(new ArrayList<>());

expectLockObtainedFor(taskId00, taskId01);
expectDirectoryNotEmpty(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 Down Expand Up @@ -2629,7 +2601,6 @@ public void shouldNotCommitNonCorruptedRestoringActiveTasksAndNotCommitRunningSt
when(tasks.allTasksPerId()).thenReturn(mkMap(mkEntry(taskId02, corruptedTask)));
when(tasks.task(taskId02)).thenReturn(corruptedTask);
final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
when(stateUpdater.getTasks()).thenReturn(mkSet(activeRestoringTask, standbyTask));

@clolov clolov Mar 18, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This was reported as unnecessary stubbing by Mockito

when(consumer.assignment()).thenReturn(intersection(HashSet::new, taskId00Partitions, taskId01Partitions, taskId02Partitions));

taskManager.handleCorruption(mkSet(taskId02));
Expand Down Expand Up @@ -2713,7 +2684,7 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
@Test
public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
final ProcessorStateManager stateManager = Mockito.mock(ProcessorStateManager.class);
expect(stateDirectory.listNonEmptyTaskDirectories()).andStubReturn(new ArrayList<>());
when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>());

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

Expand All @@ -2733,8 +2704,6 @@ public void shouldNotAttemptToCommitInHandleCorruptedDuringARebalance() {
.thenReturn(assignment)
.thenReturn(union(HashSet::new, taskId00Partitions, taskId01Partitions));

replay(stateDirectory);

uncorruptedActive.setCommittableOffsetsAndMetadata(offsets);

taskManager.handleAssignment(firstAssignement, emptyMap());
Expand Down Expand Up @@ -3382,7 +3351,7 @@ public void shouldNotCommitCreatedTasksOnRevocationOrClosure() {

taskManager.handleAssignment(emptyMap(), emptyMap());
assertThat(task00.state(), is(Task.State.CLOSED));
Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(eq(taskId00));
Mockito.verify(activeTaskCreator).closeAndRemoveTaskProducerIfNeeded(taskId00);
}

@Test
Expand Down Expand Up @@ -3882,9 +3851,9 @@ public void shouldInitialiseNewStandbyTasks() {

@Test
public void shouldHandleRebalanceEvents() {
final Set<TopicPartition> assignment = singleton(new TopicPartition("assignment", 0));
when(consumer.assignment()).thenReturn(assignment);
expect(stateDirectory.listNonEmptyTaskDirectories()).andReturn(new ArrayList<>());
replay(stateDirectory);
when(stateDirectory.listNonEmptyTaskDirectories()).thenReturn(new ArrayList<>());
assertThat(taskManager.rebalanceInProgress(), is(false));
taskManager.handleRebalanceStart(emptySet());
assertThat(taskManager.rebalanceInProgress(), is(true));
Expand Down Expand Up @@ -4000,8 +3969,6 @@ public void shouldNotCommitActiveAndStandbyTasksWhileRebalanceInProgress() throw
when(standbyTaskCreator.createTasks(taskId01Assignment))
.thenReturn(singletonList(task01));

replay(stateDirectory);

taskManager.handleAssignment(taskId00Assignment, taskId01Assignment);
assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));

Expand Down Expand Up @@ -4720,28 +4687,21 @@ private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<Top
return allTasks;
}

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

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

private void expectUnlockFor(final TaskId... tasks) throws Exception {
private void expectLockFailedFor(final TaskId... tasks) {
for (final TaskId task : tasks) {
stateDirectory.unlock(task);
expectLastCall();
when(stateDirectory.lock(task)).thenReturn(false);
}
}

private void expectDirectoryNotEmpty(final TaskId... tasks) {
for (final TaskId taskId : tasks) {
expect(stateDirectory.directoryForTaskIsEmpty(taskId)).andReturn(false);
when(stateDirectory.directoryForTaskIsEmpty(taskId)).thenReturn(false);
}
}

Expand Down Expand Up @@ -5002,14 +4962,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);
lenient().when(stateDirectory.checkpointFileFor(task)).thenReturn(checkpointFile);

@clolov clolov Jan 24, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is lenient because shouldComputeOffsetSumForRestoringActiveTaskWithStateUpdater and shouldComputeOffsetSumForRestoringStandbyTaskWithStateUpdater do not invoke this method. I am equally happy to explicitly mock the required calls only within those two tests to maintain the same strictness.

expectDirectoryNotEmpty(task);
}

Expand Down