Examples with MockStateStore org.apache.kafka.test.MockStateStore used on opensource projects

Example 6 with MockStateStore

use of org.apache.kafka.test.MockStateStore in project apache-kafka-on-k8s by banzaicloud.

the class ProcessorStateManagerTest method shouldSuccessfullyReInitializeStateStores.

@SuppressWarnings("unchecked")
@Test
public void shouldSuccessfullyReInitializeStateStores() throws IOException {
    final String store2Name = "store2";
    final String store2Changelog = "store2-changelog";
    final TopicPartition store2Partition = new TopicPartition(store2Changelog, 0);
    final List<TopicPartition> changelogPartitions = Arrays.asList(changelogTopicPartition, store2Partition);
    Map<String, String> storeToChangelog = new HashMap() {

        {
            put(storeName, changelogTopic);
            put(store2Name, store2Changelog);
        }
    };
    final ProcessorStateManager stateManager = new ProcessorStateManager(taskId, changelogPartitions, false, stateDirectory, storeToChangelog, changelogReader, false, logContext);
    final MockStateStore stateStore = new MockStateStore(storeName, true);
    final MockStateStore stateStore2 = new MockStateStore(store2Name, true);
    stateManager.register(stateStore, stateStore.stateRestoreCallback);
    stateManager.register(stateStore2, stateStore2.stateRestoreCallback);
    stateStore.initialized = false;
    stateStore2.initialized = false;
    stateManager.reinitializeStateStoresForPartitions(changelogPartitions, new NoOpProcessorContext() {

        @Override
        public void register(final StateStore store, final boolean deprecatedAndIgnoredLoggingEnabled, final StateRestoreCallback stateRestoreCallback) {
            stateManager.register(store, stateRestoreCallback);
        }
    });
    assertTrue(stateStore.initialized);
    assertTrue(stateStore2.initialized);
}

Example 7 with MockStateStore

use of org.apache.kafka.test.MockStateStore in project apache-kafka-on-k8s by banzaicloud.

the class ProcessorStateManagerTest method testGetStore.

@Test
public void testGetStore() throws IOException {
    final MockStateStore mockStateStore = new MockStateStore(nonPersistentStoreName, false);
    final ProcessorStateManager stateMgr = new ProcessorStateManager(new TaskId(0, 1), noPartitions, false, stateDirectory, Collections.<String, String>emptyMap(), changelogReader, false, logContext);
    try {
        stateMgr.register(mockStateStore, mockStateStore.stateRestoreCallback);
        assertNull(stateMgr.getStore("noSuchStore"));
        assertEquals(mockStateStore, stateMgr.getStore(nonPersistentStoreName));
    } finally {
        stateMgr.close(Collections.<TopicPartition, Long>emptyMap());
    }
}

Example 8 with MockStateStore

use of org.apache.kafka.test.MockStateStore in project apache-kafka-on-k8s by banzaicloud.

the class ProcessorStateManagerTest method testRegisterNonPersistentStore.

@Test
public void testRegisterNonPersistentStore() throws IOException {
    final MockStateStore nonPersistentStore = // non persistent store
    new MockStateStore(nonPersistentStoreName, false);
    final ProcessorStateManager stateMgr = new ProcessorStateManager(new TaskId(0, 2), noPartitions, false, stateDirectory, new HashMap<String, String>() {

        {
            put(persistentStoreName, persistentStoreTopicName);
            put(nonPersistentStoreName, nonPersistentStoreTopicName);
        }
    }, changelogReader, false, logContext);
    try {
        stateMgr.register(nonPersistentStore, nonPersistentStore.stateRestoreCallback);
        assertTrue(changelogReader.wasRegistered(new TopicPartition(nonPersistentStoreTopicName, 2)));
    } finally {
        stateMgr.close(Collections.<TopicPartition, Long>emptyMap());
    }
}

Example 9 with MockStateStore

use of org.apache.kafka.test.MockStateStore in project apache-kafka-on-k8s by banzaicloud.

the class ProcessorStateManagerTest method shouldCloseAllStoresEvenIfStoreThrowsExcepiton.

@Test
public void shouldCloseAllStoresEvenIfStoreThrowsExcepiton() throws IOException {
    final ProcessorStateManager stateManager = new ProcessorStateManager(taskId, Collections.singleton(changelogTopicPartition), false, stateDirectory, Collections.singletonMap(storeName, changelogTopic), changelogReader, false, logContext);
    final AtomicBoolean closedStore = new AtomicBoolean(false);
    final MockStateStore stateStore1 = new MockStateStore(storeName, true) {

        @Override
        public void close() {
            throw new RuntimeException("KABOOM!");
        }
    };
    final MockStateStore stateStore2 = new MockStateStore(storeName + "2", true) {

        @Override
        public void close() {
            closedStore.set(true);
        }
    };
    stateManager.register(stateStore1, stateStore1.stateRestoreCallback);
    stateManager.register(stateStore2, stateStore2.stateRestoreCallback);
    try {
        stateManager.close(Collections.<TopicPartition, Long>emptyMap());
    } catch (final ProcessorStateException expected) {
    /* ignode */
    }
    Assert.assertTrue(closedStore.get());
}

Example 10 with MockStateStore

use of org.apache.kafka.test.MockStateStore in project apache-kafka-on-k8s by banzaicloud.

the class ProcessorStateManagerTest method shouldRestoreStoreWithBatchingRestoreSpecification.

@Test
public void shouldRestoreStoreWithBatchingRestoreSpecification() throws Exception {
    final TaskId taskId = new TaskId(0, 2);
    final MockBatchingStateRestoreListener batchingRestoreCallback = new MockBatchingStateRestoreListener();
    final KeyValue<byte[], byte[]> expectedKeyValue = KeyValue.pair(key, value);
    final MockStateStore persistentStore = getPersistentStore();
    final ProcessorStateManager stateMgr = getStandByStateManager(taskId);
    try {
        stateMgr.register(persistentStore, batchingRestoreCallback);
        stateMgr.updateStandbyStates(persistentStorePartition, Collections.singletonList(consumerRecord));
        assertThat(batchingRestoreCallback.getRestoredRecords().size(), is(1));
        assertTrue(batchingRestoreCallback.getRestoredRecords().contains(expectedKeyValue));
    } finally {
        stateMgr.close(Collections.<TopicPartition, Long>emptyMap());
    }
}