TransactionalStateTaskRestoreManager implements TaskRestoreManager { @VisibleForTesting static Map<SystemStreamPartition, SystemStreamPartitionMetadata> getCurrentChangelogOffsets( TaskModel taskModel, Map<String, SystemStream> storeChangelogs, SSPMetadataCache sspMetadataCache) { Map<SystemStreamPartition, SystemStreamPartitionMetadata> changelogOffsets = new HashMap<>(); Partition changelogPartition = taskModel.getChangelogPartition(); for (Map.Entry<String, SystemStream> storeChangelog : storeChangelogs.entrySet()) { SystemStream changelog = storeChangelog.getValue(); SystemStreamPartition changelogSSP = new SystemStreamPartition( changelog.getSystem(), changelog.getStream(), changelogPartition); SystemStreamPartitionMetadata metadata = sspMetadataCache.getMetadata(changelogSSP); changelogOffsets.put(changelogSSP, metadata); } LOG.info("Got current changelog offsets for taskName: {} as: {}", taskModel.getTaskName(), changelogOffsets); return changelogOffsets; } TransactionalStateTaskRestoreManager(
      TaskModel taskModel,
      Map<String, StorageEngine> storeEngines,
      Map<String, SystemStream> storeChangelogs,
      SystemAdmins systemAdmins,
      Map<String, SystemConsumer> storeConsumers,
      SSPMetadataCache sspMetadataCache,
      File loggedStoreBaseDirectory,
      File nonLoggedStoreBaseDirectory,
      Config config,
      Clock clock); @Override void init(Map<SystemStreamPartition, String> checkpointedChangelogOffsets); @Override void restore(); void stopPersistentStores();  }
@Test public void testGetCurrentChangelogOffsets() { TaskModel mockTaskModel = mock(TaskModel.class); when(mockTaskModel.getTaskName()).thenReturn(new TaskName("Partition 0")); Partition taskChangelogPartition = new Partition(0); when(mockTaskModel.getChangelogPartition()).thenReturn(taskChangelogPartition); String store1Name = "store1"; String changelog1SystemName = "system1"; String changelog1StreamName = "store1Changelog"; String store2Name = "store2"; String changelog2SystemName = "system2"; String changelog2StreamName = "store2Changelog"; SystemStream changelog1SystemStream = new SystemStream(changelog1SystemName, changelog1StreamName); SystemStream changelog2SystemStream = new SystemStream(changelog2SystemName, changelog2StreamName); Map<String, SystemStream> mockStoreChangelogs = ImmutableMap.of( store1Name, changelog1SystemStream, store2Name, changelog2SystemStream); SSPMetadataCache mockSSPMetadataCache = mock(SSPMetadataCache.class); SystemStreamPartition changelog1SSP = new SystemStreamPartition(changelog1SystemStream, taskChangelogPartition); SystemStreamPartitionMetadata changelog1SSPMetadata = new SystemStreamPartitionMetadata("0", "1", "2"); when(mockSSPMetadataCache.getMetadata(eq(changelog1SSP))).thenReturn(changelog1SSPMetadata); SystemStreamPartition changelog2SSP = new SystemStreamPartition(changelog2SystemStream, taskChangelogPartition); SystemStreamPartitionMetadata changelog2SSPMetadata = new SystemStreamPartitionMetadata("1", "2", "3"); when(mockSSPMetadataCache.getMetadata(eq(changelog2SSP))).thenReturn(changelog2SSPMetadata); Map<SystemStreamPartition, SystemStreamPartitionMetadata> currentChangelogOffsets = TransactionalStateTaskRestoreManager.getCurrentChangelogOffsets( mockTaskModel, mockStoreChangelogs, mockSSPMetadataCache); verify(mockSSPMetadataCache, times(1)).getMetadata(changelog1SSP); verify(mockSSPMetadataCache, times(1)).getMetadata(changelog2SSP); verifyNoMoreInteractions(mockSSPMetadataCache); assertEquals(2, currentChangelogOffsets.size()); assertEquals(changelog1SSPMetadata, currentChangelogOffsets.get(changelog1SSP)); assertEquals(changelog2SSPMetadata, currentChangelogOffsets.get(changelog2SSP)); }