Examples with EasyMock.expectLastCall org.easymock.EasyMock.expectLastCall used on opensource projects

Example 1 with EasyMock.expectLastCall

use of org.easymock.EasyMock.expectLastCall in project kafka by apache.

the class TaskManagerTest method shouldReviveCorruptTasks.

@Test
public void shouldReviveCorruptTasks() {
    final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class);
    stateManager.markChangelogAsCorrupted(taskId00Partitions);
    EasyMock.expectLastCall().once();
    replay(stateManager);
    final AtomicBoolean enforcedCheckpoint = new AtomicBoolean(false);
    final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {

        @Override
        public void postCommit(final boolean enforceCheckpoint) {
            if (enforceCheckpoint) {
                enforcedCheckpoint.set(true);
            }
            super.postCommit(enforceCheckpoint);
        }
    };
    // `handleAssignment`
    expectRestoreToBeCompleted(consumer, changeLogReader);
    expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00));
    topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString());
    expectLastCall().anyTimes();
    expect(consumer.assignment()).andReturn(taskId00Partitions);
    replay(activeTaskCreator, topologyBuilder, consumer, changeLogReader);
    taskManager.handleAssignment(taskId00Assignment, emptyMap());
    assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
    assertThat(task00.state(), is(Task.State.RUNNING));
    task00.setChangelogOffsets(singletonMap(t1p0, 0L));
    taskManager.handleCorruption(singleton(taskId00));
    assertThat(task00.commitPrepared, is(true));
    assertThat(task00.state(), is(Task.State.CREATED));
    assertThat(task00.partitionsForOffsetReset, equalTo(taskId00Partitions));
    assertThat(enforcedCheckpoint.get(), is(true));
    assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
    assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
    verify(stateManager);
    verify(consumer);
}

Example 2 with EasyMock.expectLastCall

use of org.easymock.EasyMock.expectLastCall in project kafka by apache.

the class StoreChangelogReaderTest method shouldRequestCommittedOffsetsAndHandleTimeoutException.

@Test
public void shouldRequestCommittedOffsetsAndHandleTimeoutException() {
    final TaskId taskId = new TaskId(0, 0);
    final Task mockTask = mock(Task.class);
    if (type == ACTIVE) {
        mockTask.clearTaskTimeout();
    }
    mockTask.maybeInitTaskTimeoutOrThrow(anyLong(), anyObject());
    EasyMock.expectLastCall();
    EasyMock.expect(stateManager.changelogAsSource(tp)).andReturn(true).anyTimes();
    EasyMock.expect(storeMetadata.offset()).andReturn(5L).anyTimes();
    EasyMock.expect(stateManager.changelogOffsets()).andReturn(singletonMap(tp, 5L));
    EasyMock.expect(stateManager.taskId()).andReturn(taskId).anyTimes();
    EasyMock.replay(mockTask, stateManager, storeMetadata, store);
    final AtomicBoolean functionCalled = new AtomicBoolean(false);
    final MockConsumer<byte[], byte[]> consumer = new MockConsumer<byte[], byte[]>(OffsetResetStrategy.EARLIEST) {

        @Override
        public Map<TopicPartition, OffsetAndMetadata> committed(final Set<TopicPartition> partitions) {
            if (functionCalled.get()) {
                return partitions.stream().collect(Collectors.toMap(Function.identity(), partition -> new OffsetAndMetadata(10L)));
            } else {
                functionCalled.set(true);
                throw new TimeoutException("KABOOM!");
            }
        }
    };
    adminClient.updateEndOffsets(Collections.singletonMap(tp, 20L));
    final StoreChangelogReader changelogReader = new StoreChangelogReader(time, config, logContext, adminClient, consumer, callback);
    changelogReader.setMainConsumer(consumer);
    changelogReader.register(tp, stateManager);
    changelogReader.restore(Collections.singletonMap(taskId, mockTask));
    assertEquals(type == ACTIVE ? StoreChangelogReader.ChangelogState.REGISTERED : StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
    if (type == ACTIVE) {
        assertNull(changelogReader.changelogMetadata(tp).endOffset());
    } else {
        assertEquals(0L, (long) changelogReader.changelogMetadata(tp).endOffset());
    }
    assertTrue(functionCalled.get());
    verify(mockTask);
    resetToDefault(mockTask);
    if (type == ACTIVE) {
        mockTask.clearTaskTimeout();
        mockTask.clearTaskTimeout();
        expectLastCall();
    }
    replay(mockTask);
    changelogReader.restore(Collections.singletonMap(taskId, mockTask));
    assertEquals(StoreChangelogReader.ChangelogState.RESTORING, changelogReader.changelogMetadata(tp).state());
    assertEquals(type == ACTIVE ? 10L : 0L, (long) changelogReader.changelogMetadata(tp).endOffset());
    assertEquals(6L, consumer.position(tp));
    verify(mockTask);
}