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

Example 1 with MockRestoreCallback

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

the class StoreChangelogReaderTest method shouldRestoreMultipleStores.

@Test
public void shouldRestoreMultipleStores() {
    final TopicPartition one = new TopicPartition("one", 0);
    final TopicPartition two = new TopicPartition("two", 0);
    final MockRestoreCallback callbackOne = new MockRestoreCallback();
    final MockRestoreCallback callbackTwo = new MockRestoreCallback();
    final CompositeRestoreListener restoreListener1 = new CompositeRestoreListener(callbackOne);
    final CompositeRestoreListener restoreListener2 = new CompositeRestoreListener(callbackTwo);
    setupConsumer(10, topicPartition);
    setupConsumer(5, one);
    setupConsumer(3, two);
    changelogReader.register(new StateRestorer(topicPartition, restoreListener, null, Long.MAX_VALUE, true, "storeName1"));
    changelogReader.register(new StateRestorer(one, restoreListener1, null, Long.MAX_VALUE, true, "storeName2"));
    changelogReader.register(new StateRestorer(two, restoreListener2, null, Long.MAX_VALUE, true, "storeName3"));
    expect(active.restoringTaskFor(one)).andReturn(null);
    expect(active.restoringTaskFor(two)).andReturn(null);
    replay(active);
    changelogReader.restore(active);
    assertThat(callback.restored.size(), equalTo(10));
    assertThat(callbackOne.restored.size(), equalTo(5));
    assertThat(callbackTwo.restored.size(), equalTo(3));
}

Example 2 with MockRestoreCallback

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

the class StoreChangelogReaderTest method shouldRestorePartitionsRegisteredPostInitialization.

@Test
public void shouldRestorePartitionsRegisteredPostInitialization() {
    final MockRestoreCallback callbackTwo = new MockRestoreCallback();
    final CompositeRestoreListener restoreListener2 = new CompositeRestoreListener(callbackTwo);
    setupConsumer(1, topicPartition);
    consumer.updateEndOffsets(Collections.singletonMap(topicPartition, 10L));
    changelogReader.register(new StateRestorer(topicPartition, restoreListener, null, Long.MAX_VALUE, false, "storeName"));
    final TopicPartition postInitialization = new TopicPartition("other", 0);
    expect(active.restoringTaskFor(topicPartition)).andReturn(null);
    expect(active.restoringTaskFor(topicPartition)).andReturn(null);
    expect(active.restoringTaskFor(postInitialization)).andReturn(null);
    replay(active);
    assertTrue(changelogReader.restore(active).isEmpty());
    addRecords(9, topicPartition, 1);
    setupConsumer(3, postInitialization);
    consumer.updateBeginningOffsets(Collections.singletonMap(postInitialization, 0L));
    consumer.updateEndOffsets(Collections.singletonMap(postInitialization, 3L));
    changelogReader.register(new StateRestorer(postInitialization, restoreListener2, null, Long.MAX_VALUE, false, "otherStore"));
    final Collection<TopicPartition> expected = Utils.mkSet(topicPartition, postInitialization);
    consumer.assign(expected);
    assertThat(changelogReader.restore(active), equalTo(expected));
    assertThat(callback.restored.size(), equalTo(10));
    assertThat(callbackTwo.restored.size(), equalTo(3));
}

Example 3 with MockRestoreCallback

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

the class AbstractTaskTest method shouldDeleteAndRecreateStoreDirectoryOnReinitialize.

@Test
public void shouldDeleteAndRecreateStoreDirectoryOnReinitialize() throws IOException {
    final StreamsConfig streamsConfig = new StreamsConfig(new Properties() {

        {
            put(StreamsConfig.APPLICATION_ID_CONFIG, "app-id");
            put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
            put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getAbsolutePath());
        }
    });
    final Consumer consumer = EasyMock.createNiceMock(Consumer.class);
    final StateStore store1 = EasyMock.createNiceMock(StateStore.class);
    final StateStore store2 = EasyMock.createNiceMock(StateStore.class);
    final StateStore store3 = EasyMock.createNiceMock(StateStore.class);
    final StateStore store4 = EasyMock.createNiceMock(StateStore.class);
    final String storeName1 = "storeName1";
    final String storeName2 = "storeName2";
    final String storeName3 = "storeName3";
    final String storeName4 = "storeName4";
    expect(store1.name()).andReturn(storeName1).anyTimes();
    EasyMock.replay(store1);
    expect(store2.name()).andReturn(storeName2).anyTimes();
    EasyMock.replay(store2);
    expect(store3.name()).andReturn(storeName3).anyTimes();
    EasyMock.replay(store3);
    expect(store4.name()).andReturn(storeName4).anyTimes();
    EasyMock.replay(store4);
    final StateDirectory stateDirectory = new StateDirectory(streamsConfig, new MockTime());
    final AbstractTask task = createTask(consumer, new HashMap<StateStore, String>() {

        {
            put(store1, storeTopicPartition1.topic());
            put(store2, storeTopicPartition2.topic());
            put(store3, storeTopicPartition3.topic());
            put(store4, storeTopicPartition4.topic());
        }
    }, stateDirectory);
    final String taskDir = stateDirectory.directoryForTask(task.id).getAbsolutePath();
    final File storeDirectory1 = new File(taskDir + File.separator + "rocksdb" + File.separator + storeName1);
    final File storeDirectory2 = new File(taskDir + File.separator + "rocksdb" + File.separator + storeName2);
    final File storeDirectory3 = new File(taskDir + File.separator + storeName3);
    final File storeDirectory4 = new File(taskDir + File.separator + storeName4);
    final File testFile1 = new File(storeDirectory1.getAbsolutePath() + File.separator + "testFile");
    final File testFile2 = new File(storeDirectory2.getAbsolutePath() + File.separator + "testFile");
    final File testFile3 = new File(storeDirectory3.getAbsolutePath() + File.separator + "testFile");
    final File testFile4 = new File(storeDirectory4.getAbsolutePath() + File.separator + "testFile");
    storeDirectory1.mkdirs();
    storeDirectory2.mkdirs();
    storeDirectory3.mkdirs();
    storeDirectory4.mkdirs();
    testFile1.createNewFile();
    assertTrue(testFile1.exists());
    testFile2.createNewFile();
    assertTrue(testFile2.exists());
    testFile3.createNewFile();
    assertTrue(testFile3.exists());
    testFile4.createNewFile();
    assertTrue(testFile4.exists());
    task.processorContext = new InternalMockProcessorContext(stateDirectory.directoryForTask(task.id), streamsConfig);
    task.stateMgr.register(store1, new MockRestoreCallback());
    task.stateMgr.register(store2, new MockRestoreCallback());
    task.stateMgr.register(store3, new MockRestoreCallback());
    task.stateMgr.register(store4, new MockRestoreCallback());
    // only reinitialize store1 and store3 -- store2 and store4 should be untouched
    task.reinitializeStateStoresForPartitions(Utils.mkSet(storeTopicPartition1, storeTopicPartition3));
    assertFalse(testFile1.exists());
    assertTrue(testFile2.exists());
    assertFalse(testFile3.exists());
    assertTrue(testFile4.exists());
}

Example 4 with MockRestoreCallback

use of org.apache.kafka.test.MockRestoreCallback in project kafka by apache.

the class StoreChangelogReaderTest method shouldRestoreMultipleStores.

@Test
public void shouldRestoreMultipleStores() throws Exception {
    final TopicPartition one = new TopicPartition("one", 0);
    final TopicPartition two = new TopicPartition("two", 0);
    final MockRestoreCallback callbackOne = new MockRestoreCallback();
    final MockRestoreCallback callbackTwo = new MockRestoreCallback();
    setupConsumer(10, topicPartition);
    setupConsumer(5, one);
    setupConsumer(3, two);
    changelogReader.register(new StateRestorer(topicPartition, callback, null, Long.MAX_VALUE, true));
    changelogReader.register(new StateRestorer(one, callbackOne, null, Long.MAX_VALUE, true));
    changelogReader.register(new StateRestorer(two, callbackTwo, null, Long.MAX_VALUE, true));
    changelogReader.restore();
    assertThat(callback.restoreCount, equalTo(10));
    assertThat(callbackOne.restoreCount, equalTo(5));
    assertThat(callbackTwo.restoreCount, equalTo(3));
}