Examples with CheckpointId org.apache.samza.checkpoint.CheckpointId used on opensource projects

Example 41 with CheckpointId

use of org.apache.samza.checkpoint.CheckpointId in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testGetStoreActionsForLoggedPersistentStore_RestoreFromLocalToNewestIfCheckpointedOffsetInRangeButMaybeCompacted.

/**
 * This can happen if the changelog offset is valid but the checkpoint is older than min compaction lag ms. E.g., when
 * the job/container shut down and restarted after a long time.
 */
@Test
public void testGetStoreActionsForLoggedPersistentStore_RestoreFromLocalToNewestIfCheckpointedOffsetInRangeButMaybeCompacted() {
    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);
    Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of(store1Name, store1Engine);
    String changelog1SystemName = "system1";
    String changelog1StreamName = "store1Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    // checkpointed changelog offset is valid
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("4", "20", "21");
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream);
    String changelog1CheckpointedOffset = "5";
    // checkpoint timestamp older than default min compaction lag
    CheckpointId checkpointId = CheckpointId.deserialize("0-0");
    KafkaStateCheckpointMarker kafkaStateCheckpointMarker = new KafkaStateCheckpointMarker(changelog1SSP, changelog1CheckpointedOffset);
    Map<String, KafkaStateCheckpointMarker> mockCheckpointedChangelogOffset = new HashMap<String, KafkaStateCheckpointMarker>() {

        {
            put(store1Name, kafkaStateCheckpointMarker);
        }
    };
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogOffsets = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(changelog1SSP.getSystem())).thenReturn(mockSystemAdmin);
    StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class);
    File mockLoggedStoreBaseDir = mock(File.class);
    File mockNonLoggedStoreBaseDir = mock(File.class);
    Config mockConfig = mock(Config.class);
    Clock mockClock = mock(Clock.class);
    File mockCurrentStoreDir = mock(File.class);
    File mockStoreNewerCheckpointDir = mock(File.class);
    File mockStoreOlderCheckpointDir = mock(File.class);
    String olderCheckpointDirLocalOffset = "3";
    String newerCheckpointDirLocalOffset = "5";
    when(mockStorageManagerUtil.getTaskStoreDir(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(mockCurrentStoreDir);
    when(mockStorageManagerUtil.getTaskStoreCheckpointDirs(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(ImmutableList.of(mockStoreNewerCheckpointDir, mockStoreOlderCheckpointDir));
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreNewerCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreOlderCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    Set<SystemStreamPartition> mockChangelogSSPs = ImmutableSet.of(changelog1SSP);
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreNewerCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(ImmutableMap.of(changelog1SSP, newerCheckpointDirLocalOffset));
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreOlderCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(// less than checkpointed offset (5)
    ImmutableMap.of(changelog1SSP, olderCheckpointDirLocalOffset));
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    StoreActions storeActions = TransactionalStateTaskRestoreManager.getStoreActions(mockTaskModel, mockStoreEngines, mockStoreChangelogs, mockCheckpointedChangelogOffset, checkpointId, mockCurrentChangelogOffsets, mockSystemAdmins, mockStorageManagerUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir, mockConfig, mockClock);
    // ensure that the current store dir and older checkpoint dir are marked for deletion
    assertEquals(2, storeActions.storeDirsToDelete.get(store1Name).size());
    assertTrue(storeActions.storeDirsToDelete.get(store1Name).contains(mockCurrentStoreDir));
    assertTrue(storeActions.storeDirsToDelete.get(store1Name).contains(mockStoreOlderCheckpointDir));
    // ensure that newer checkpoint dir is retained
    assertEquals(1, storeActions.storeDirsToRetain.size());
    assertEquals(mockStoreNewerCheckpointDir, storeActions.storeDirsToRetain.get(store1Name));
    // ensure that we mark the store for restore to head (from local checkpoint to current newest)
    assertEquals("5", storeActions.storesToRestore.get(store1Name).startingOffset);
    assertEquals("20", storeActions.storesToRestore.get(store1Name).endingOffset);
}

Example 42 with CheckpointId

use of org.apache.samza.checkpoint.CheckpointId in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testGetStoreActionsForLoggedPersistentStore_NoRestoreButTrimIfUpToDateStoreCheckpoint.

/**
 * Ensure that we do a trim even if the local offset == checkpointed changelog offset, since there may be
 * additional messages in the changelog since the last commit that we need to revert.
 */
@Test
public void testGetStoreActionsForLoggedPersistentStore_NoRestoreButTrimIfUpToDateStoreCheckpoint() {
    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);
    Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of(store1Name, store1Engine);
    String changelog1SystemName = "system1";
    String changelog1StreamName = "store1Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "10", "11");
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream);
    String changelog1CheckpointedOffset = "5";
    CheckpointId checkpointId = CheckpointId.create();
    KafkaStateCheckpointMarker kafkaStateCheckpointMarker = new KafkaStateCheckpointMarker(changelog1SSP, changelog1CheckpointedOffset);
    ImmutableMap<String, KafkaStateCheckpointMarker> mockCheckpointedChangelogOffset = ImmutableMap.of(store1Name, kafkaStateCheckpointMarker);
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogOffsets = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(changelog1SSP.getSystem())).thenReturn(mockSystemAdmin);
    StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class);
    File mockLoggedStoreBaseDir = mock(File.class);
    File mockNonLoggedStoreBaseDir = mock(File.class);
    Config mockConfig = mock(Config.class);
    Clock mockClock = mock(Clock.class);
    File mockCurrentStoreDir = mock(File.class);
    File mockStoreUpToDateCheckpointDir = mock(File.class);
    File mockStoreOlderCheckpointDir = mock(File.class);
    when(mockStorageManagerUtil.getTaskStoreDir(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(mockCurrentStoreDir);
    when(mockStorageManagerUtil.getTaskStoreCheckpointDirs(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(ImmutableList.of(mockStoreUpToDateCheckpointDir, mockStoreOlderCheckpointDir));
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreUpToDateCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreOlderCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    Set<SystemStreamPartition> mockChangelogSSPs = ImmutableSet.of(changelog1SSP);
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreUpToDateCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(ImmutableMap.of(changelog1SSP, changelog1CheckpointedOffset));
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreOlderCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(// less than checkpointed offset (5)
    ImmutableMap.of(changelog1SSP, "3"));
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    StoreActions storeActions = TransactionalStateTaskRestoreManager.getStoreActions(mockTaskModel, mockStoreEngines, mockStoreChangelogs, mockCheckpointedChangelogOffset, checkpointId, mockCurrentChangelogOffsets, mockSystemAdmins, mockStorageManagerUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir, mockConfig, mockClock);
    // ensure that both the current dir and older checkpoint dir are marked for deletion
    assertEquals(2, storeActions.storeDirsToDelete.get(store1Name).size());
    assertTrue(storeActions.storeDirsToDelete.get(store1Name).contains(mockCurrentStoreDir));
    assertTrue(storeActions.storeDirsToDelete.get(store1Name).contains(mockStoreOlderCheckpointDir));
    // ensure that the up-to-date store checkpoint is marked for retention
    assertEquals(mockStoreUpToDateCheckpointDir, storeActions.storeDirsToRetain.get(store1Name));
    // ensure that we mark the store for restore even if local offset == checkpointed offset
    // this is required even if there are no messages to restore, since there may be message we need to trim
    assertEquals(changelog1CheckpointedOffset, storeActions.storesToRestore.get(store1Name).startingOffset);
    assertEquals(changelog1CheckpointedOffset, storeActions.storesToRestore.get(store1Name).endingOffset);
}

Example 43 with CheckpointId

use of org.apache.samza.checkpoint.CheckpointId in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testGetStoreActionsForLoggedNonPersistentStore_FullRestoreIfCheckpointedOffsetNewerThanNewest.

/**
 * This can happen if the changelog topic was manually deleted and recreated, and the checkpointed/local changelog
 * offset is not valid anymore.
 */
@Test
public void testGetStoreActionsForLoggedNonPersistentStore_FullRestoreIfCheckpointedOffsetNewerThanNewest() {
    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);
    // non-persistent store
    when(mockStore1Properties.isPersistedToDisk()).thenReturn(false);
    Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of(store1Name, store1Engine);
    String changelog1SystemName = "system1";
    String changelog1StreamName = "store1Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    // checkpointed changelog offset > newest offset (e.g. changelog topic got changed)
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "10", "11");
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream);
    String changelog1CheckpointedOffset = "21";
    CheckpointId checkpointId = CheckpointId.create();
    KafkaStateCheckpointMarker kafkaStateCheckpointMarker = new KafkaStateCheckpointMarker(changelog1SSP, changelog1CheckpointedOffset);
    Map<String, KafkaStateCheckpointMarker> mockCheckpointedChangelogOffset = new HashMap<String, KafkaStateCheckpointMarker>() {

        {
            put(store1Name, kafkaStateCheckpointMarker);
        }
    };
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogOffsets = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(changelog1SSP.getSystem())).thenReturn(mockSystemAdmin);
    StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class);
    File mockLoggedStoreBaseDir = mock(File.class);
    File mockNonLoggedStoreBaseDir = mock(File.class);
    Config mockConfig = mock(Config.class);
    Clock mockClock = mock(Clock.class);
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    StoreActions storeActions = TransactionalStateTaskRestoreManager.getStoreActions(mockTaskModel, mockStoreEngines, mockStoreChangelogs, mockCheckpointedChangelogOffset, checkpointId, mockCurrentChangelogOffsets, mockSystemAdmins, mockStorageManagerUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir, mockConfig, mockClock);
    // ensure that there is nothing to retain or delete
    assertEquals(0, storeActions.storeDirsToDelete.size());
    assertEquals(0, storeActions.storeDirsToRetain.size());
    // ensure that we mark the store for full restore (from current oldest to current newest)
    assertEquals("0", storeActions.storesToRestore.get(store1Name).startingOffset);
    assertEquals("10", storeActions.storesToRestore.get(store1Name).endingOffset);
}

Example 44 with CheckpointId

use of org.apache.samza.checkpoint.CheckpointId in project samza by apache.

the class TestTransactionalStateTaskRestoreManager method testGetStoreActionsForLoggedPersistentStore_RetainOneCheckpointIfMultipleCheckpointsWithSameOffset.

/**
 * This can happen if no new messages were written to the store between commits. There may be more than one store
 * checkpoint if container fails during commit after creating a checkpoint but before deleting the old one.
 */
@Test
public void testGetStoreActionsForLoggedPersistentStore_RetainOneCheckpointIfMultipleCheckpointsWithSameOffset() {
    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);
    Map<String, StorageEngine> mockStoreEngines = ImmutableMap.of(store1Name, store1Engine);
    String changelog1SystemName = "system1";
    String changelog1StreamName = "store1Changelog";
    SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName);
    SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition);
    SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "10", "11");
    Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of(store1Name, changelog1SystemStream);
    String changelog1CheckpointedOffset = "5";
    CheckpointId checkpointId = CheckpointId.create();
    KafkaStateCheckpointMarker kafkaStateCheckpointMarker = new KafkaStateCheckpointMarker(changelog1SSP, changelog1CheckpointedOffset);
    ImmutableMap<String, KafkaStateCheckpointMarker> mockCheckpointedChangelogOffset = ImmutableMap.of(store1Name, kafkaStateCheckpointMarker);
    Map<SystemStreamPartition, SystemStreamPartitionMetadata> mockCurrentChangelogOffsets = ImmutableMap.of(changelog1SSP, changelog1SSPMetadata);
    SystemAdmins mockSystemAdmins = mock(SystemAdmins.class);
    SystemAdmin mockSystemAdmin = mock(SystemAdmin.class);
    when(mockSystemAdmins.getSystemAdmin(changelog1SSP.getSystem())).thenReturn(mockSystemAdmin);
    StorageManagerUtil mockStorageManagerUtil = mock(StorageManagerUtil.class);
    File mockLoggedStoreBaseDir = mock(File.class);
    File mockNonLoggedStoreBaseDir = mock(File.class);
    Config mockConfig = mock(Config.class);
    Clock mockClock = mock(Clock.class);
    File mockCurrentStoreDir = mock(File.class);
    File mockStoreNewerCheckpointDir = mock(File.class);
    File mockStoreOlderCheckpointDir = mock(File.class);
    when(mockStorageManagerUtil.getTaskStoreDir(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(mockCurrentStoreDir);
    when(mockStorageManagerUtil.getTaskStoreCheckpointDirs(eq(mockLoggedStoreBaseDir), eq(store1Name), eq(taskName), any())).thenReturn(ImmutableList.of(mockStoreNewerCheckpointDir, mockStoreOlderCheckpointDir));
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreNewerCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    when(mockStorageManagerUtil.isLoggedStoreValid(eq(store1Name), eq(mockStoreOlderCheckpointDir), any(), eq(mockStoreChangelogs), eq(mockTaskModel), any(), eq(mockStoreEngines))).thenReturn(true);
    Set<SystemStreamPartition> mockChangelogSSPs = ImmutableSet.of(changelog1SSP);
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreNewerCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(// equal to checkpointed offset (5)
    ImmutableMap.of(changelog1SSP, changelog1CheckpointedOffset));
    when(mockStorageManagerUtil.readOffsetFile(eq(mockStoreOlderCheckpointDir), eq(mockChangelogSSPs), eq(false))).thenReturn(// also equal to checkpointed offset (5)
    ImmutableMap.of(changelog1SSP, changelog1CheckpointedOffset));
    Mockito.when(mockSystemAdmin.offsetComparator(anyString(), anyString())).thenAnswer((Answer<Integer>) invocation -> {
        String offset1 = (String) invocation.getArguments()[0];
        String offset2 = (String) invocation.getArguments()[1];
        return Long.valueOf(offset1).compareTo(Long.valueOf(offset2));
    });
    StoreActions storeActions = TransactionalStateTaskRestoreManager.getStoreActions(mockTaskModel, mockStoreEngines, mockStoreChangelogs, mockCheckpointedChangelogOffset, checkpointId, mockCurrentChangelogOffsets, mockSystemAdmins, mockStorageManagerUtil, mockLoggedStoreBaseDir, mockNonLoggedStoreBaseDir, mockConfig, mockClock);
    // ensure that both the current dir and one of the checkpoint dirs are marked for deletion
    assertEquals(2, storeActions.storeDirsToDelete.get(store1Name).size());
    assertTrue(storeActions.storeDirsToDelete.get(store1Name).contains(mockCurrentStoreDir));
    // ensure that the one of the store checkpoint is marked for retention
    assertNotNull(storeActions.storeDirsToRetain.get(store1Name));
    // ensure that we mark the store for restore even if local offset == checkpointed offset
    // this is required since there may be message we need to trim
    assertEquals(changelog1CheckpointedOffset, storeActions.storesToRestore.get(store1Name).startingOffset);
    assertEquals(changelog1CheckpointedOffset, storeActions.storesToRestore.get(store1Name).endingOffset);
}

Example 45 with CheckpointId

use of org.apache.samza.checkpoint.CheckpointId in project samza by apache.

the class TestBlobStoreRestoreManager method testRestoreRetainsCheckpointDirsIfValid.

@Test
public void testRestoreRetainsCheckpointDirsIfValid() throws IOException {
    String jobName = "testJobName";
    String jobId = "testJobId";
    TaskName taskName = mock(TaskName.class);
    BlobStoreRestoreManagerMetrics metrics = new BlobStoreRestoreManagerMetrics(new MetricsRegistryMap());
    metrics.initStoreMetrics(ImmutableList.of("storeName"));
    Set<String> storesToRestore = ImmutableSet.of("storeName");
    SnapshotIndex snapshotIndex = mock(SnapshotIndex.class);
    Map<String, Pair<String, SnapshotIndex>> prevStoreSnapshotIndexes = ImmutableMap.of("storeName", Pair.of("blobId", snapshotIndex));
    DirIndex dirIndex = BlobStoreTestUtil.createDirIndex("[a]");
    when(snapshotIndex.getDirIndex()).thenReturn(dirIndex);
    CheckpointId checkpointId = CheckpointId.create();
    when(snapshotIndex.getSnapshotMetadata()).thenReturn(new SnapshotMetadata(checkpointId, "jobName", "jobId", "taskName", "storeName"));
    Path loggedBaseDir = Files.createTempDirectory(BlobStoreTestUtil.TEMP_DIR_PREFIX);
    // create store dir to be deleted during restore
    Path storeDir = Files.createTempDirectory(loggedBaseDir, "storeDir-");
    // create checkpoint dir so that shouldRestore = false (areSameDir == true later)
    Path storeCheckpointDir = Files.createTempDirectory(loggedBaseDir, "storeDir-" + checkpointId + "-");
    // create a dummy file to verify after dir rename.
    Path tempFile = Files.createTempFile(storeCheckpointDir, "tempFile-", null);
    StorageConfig storageConfig = mock(StorageConfig.class);
    StorageManagerUtil storageManagerUtil = mock(StorageManagerUtil.class);
    when(storageManagerUtil.getTaskStoreDir(eq(loggedBaseDir.toFile()), eq("storeName"), eq(taskName), eq(TaskMode.Active))).thenReturn(storeDir.toFile());
    when(storageManagerUtil.getStoreCheckpointDir(any(File.class), eq(checkpointId))).thenReturn(storeCheckpointDir.toString());
    when(storageManagerUtil.getTaskStoreCheckpointDirs(any(File.class), anyString(), any(TaskName.class), any(TaskMode.class))).thenReturn(ImmutableList.of(storeCheckpointDir.toFile()));
    BlobStoreUtil blobStoreUtil = mock(BlobStoreUtil.class);
    DirDiffUtil dirDiffUtil = mock(DirDiffUtil.class);
    // ensures shouldRestore is not called
    when(dirDiffUtil.areSameDir(anySet(), anyBoolean())).thenReturn((arg1, arg2) -> true);
    // return immediately without restoring.
    when(blobStoreUtil.restoreDir(eq(storeDir.toFile()), eq(dirIndex), any(Metadata.class))).thenReturn(CompletableFuture.completedFuture(null));
    BlobStoreRestoreManager.restoreStores(jobName, jobId, taskName, storesToRestore, prevStoreSnapshotIndexes, loggedBaseDir.toFile(), storageConfig, metrics, storageManagerUtil, blobStoreUtil, dirDiffUtil, EXECUTOR);
    // verify that the store directory restore was not called (should have restored from checkpoint dir)
    verify(blobStoreUtil, times(0)).restoreDir(eq(storeDir.toFile()), eq(dirIndex), any(Metadata.class));
    // verify that the checkpoint dir was renamed to store dir
    assertFalse(storeCheckpointDir.toFile().exists());
    assertTrue(storeDir.toFile().exists());
    assertTrue(Files.exists(Paths.get(storeDir.toString(), tempFile.getFileName().toString())));
}