@Test
  public void testSetupStoreDirs() {
    TaskModel mockTaskModel = mock(TaskModel.class);
    TaskName taskName = new TaskName("Partition 0");
    when(mockTaskModel.getTaskName()).thenReturn(taskName);
    Partition taskChangelogPartition = new Partition(0);
    when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition);

    String store1Name = "store1";
    StorageEngine store1Engine = mock(StorageEngine.class);
    StoreProperties mockStore1Properties = mock(StoreProperties.class);
    when(store1Engine.getStoreProperties()).thenReturn(mockStore1Properties);
    when(mockStore1Properties.isLoggedStore()).thenReturn(true);
    when(mockStore1Properties.isPersistedToDisk()).thenReturn(true);
    String store2Name = "store2";
    StorageEngine store2Engine = mock(StorageEngine.class);
    StoreProperties mockStore2Properties = mock(StoreProperties.class);
    when(store2Engine.getStoreProperties()).thenReturn(mockStore2Properties);
    when(mockStore2Properties.isLoggedStore()).thenReturn(false); // non-logged store
    when(mockStore2Properties.isPersistedToDisk()).thenReturn(true);

    Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of(
        store1Name, store1Engine,
        store2Name, store2Engine);

    File mockStore1DirToRetain = mock(File.class);
    // there will be no dir to retain for non-logged persistent stores
    ImmutableMap<String, File> storeDirsToRetain = ImmutableMap.of(store1Name, mockStore1DirToRetain);
    ListMultimap<String, File> storeDirsToDelete = ArrayListMultimap.create();
    File mockStore1CurrentDir = mock(File.class);
    Path mockStore1CurrentDirPath = mock(Path.class);
    when(mockStore1CurrentDir.toPath()).thenReturn(mockStore1CurrentDirPath);
    File mockStore2CurrentDir = mock(File.class);
    Path mockStore2CurrentDirPath = mock(Path.class);
    when(mockStore2CurrentDir.toPath()).thenReturn(mockStore2CurrentDirPath);
    File mockStore1DirToDelete1 = mock(File.class);
    File mockStore1DirToDelete2 = mock(File.class);
    File mockStore2DirToDelete1 = mock(File.class);
    File mockStore2DirToDelete2 = mock(File.class);
    storeDirsToDelete.put(store1Name, mockStore1CurrentDir);
    storeDirsToDelete.put(store1Name, mockStore1DirToDelete1);
    storeDirsToDelete.put(store1Name, mockStore1DirToDelete2);
    storeDirsToDelete.put(store2Name, mockStore2CurrentDir);
    storeDirsToDelete.put(store2Name, mockStore2DirToDelete1);
    storeDirsToDelete.put(store2Name, mockStore2DirToDelete2);

    StoreActions storeActions = new StoreActions(storeDirsToRetain, storeDirsToDelete, ImmutableMap.of());
    StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class);
    FileUtil mockFileUtil = mock(FileUtil.class);
    File mockLoggedStoreBaseDir = mock(File.class);
    File mockNonLoggedStoreBaseDir = mock(File.class);

    when(mockStorageManagerUtil.getTaskStoreDir(eq(mockLoggedStoreBaseDir), eq(store1Name), any(), any()))
        .thenReturn(mockStore1CurrentDir);
    when(mockStorageManagerUtil.getTaskStoreDir(eq(mockNonLoggedStoreBaseDir), eq(store2Name), any(), any()))
        .thenReturn(mockStore2CurrentDir);
    when(mockFileUtil.exists(eq(mockStore1CurrentDirPath)))
        .thenReturn(false);
    when(mockFileUtil.exists(eq(mockStore2CurrentDirPath)))
        .thenReturn(true); // will not be true in reality since current dir is always cleared, but return true for testing

    TransactionalStateTaskRestoreManager.setupStoreDirs(mockTaskModel, mockStoreEngines, storeActions,
        mockStorageManagerUtil, mockFileUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir);

    // verify that store directories to delete are deleted
    verify(mockFileUtil, times(1)).rm(mockStore1CurrentDir);
    verify(mockFileUtil, times(1)).rm(mockStore1DirToDelete1);
    verify(mockFileUtil, times(1)).rm(mockStore1DirToDelete2);
    verify(mockFileUtil, times(1)).rm(mockStore2CurrentDir);
    verify(mockFileUtil, times(1)).rm(mockStore2DirToDelete1);
    verify(mockFileUtil, times(1)).rm(mockStore2DirToDelete2);

    // verify that store checkpoint directories to retain are moved to (empty) current dirs only for store 1
    // setupStoreDirs doesn't guarantee that the dir is empty by itself, but the dir will be part of dirs to delete.
    verify(mockStorageManagerUtil, times(1)).restoreCheckpointFiles(any(), any());
    verify(mockFileUtil, times(1)).exists(mockStore1CurrentDirPath);
    verify(mockFileUtil, times(1)).createDirectories(mockStore1CurrentDirPath);
    verify(mockFileUtil, times(1)).exists(mockStore2CurrentDirPath);
    verify(mockFileUtil, never()).createDirectories(mockStore2CurrentDirPath); // should not be called since exists == true

    verifyNoMoreInteractions(mockFileUtil);
  }
This can happen if the changelog topic was manually deleted and recreated, and the checkpointed/local changelog
offset is not valid anymore.
samza-core/src/test/java/org/apache/samza/storage/TestTransactionalStateTaskRestoreManager.java