package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.git.hook.GitHookService;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.junit.TemporaryDirectory;
import com.atlassian.bitbucket.mesh.partition.Partition;
import com.atlassian.bitbucket.mesh.util.Version;
import com.google.common.collect.ImmutableSet;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.nio.file.attribute.BasicFileAttributeView;
import java.nio.file.attribute.FileTime;
import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static com.atlassian.bitbucket.mesh.repository.DefaultSharedObjectsRepository.SHARED_OBJECTS;
import static com.atlassian.bitbucket.mesh.repository.ManagedDataStore.LAYOUT_VERSION;
import static com.atlassian.bitbucket.mesh.repository.ManagedDataStore.TYPE;
import static com.atlassian.bitbucket.mesh.repository.ShardState.AVAILABLE;
import static com.atlassian.bitbucket.mesh.repository.ShardState.STARTING;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.junit.Assume.assumeFalse;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;

public class ManagedShardTest {

    private static final Version BUILD_VERSION = new Version(1);
    private static final String HIERARCHY_ID = "12345123451234512345";
    private static final String PARTITION_ID = "0001";
    private static final String PARTITION_ID_LONG = "f0000000000000f";

    @Rule
    public final LogListener logListener = new LogListener(ManagedShard.class);
    @Rule
    public final TemporaryDirectory tempDir = new TemporaryDirectory();

    private Path rootDir;
    private ManagedShard shard;

    @Before
    public void setup() throws IOException {
        rootDir = tempDir.newFolder("store");

        shard = new ManagedShard(rootDir);
    }

    @Test
    public void testAttachStorage() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertSame(StorageStatus.ATTACHED, shard.attachStorage(repository));

        assertEquals(repositoryDir.resolveSibling("1.deleted"), repository.getDeletedMarker());
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "hooks"), repository.getHooksDir());
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID, HIERARCHY_ID, "lfs"), repository.getLfsDir());
        assertEquals(repositoryDir, repository.getPath());
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "templates"), repository.getTemplateDir());
    }

    @Test
    public void testAttachStorageLong() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertSame(StorageStatus.ATTACHED, shard.attachStorage(repository));

        assertEquals(repositoryDir.resolveSibling("1.deleted"), repository.getDeletedMarker());
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "hooks"), repository.getHooksDir());
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID_LONG, HIERARCHY_ID, "lfs"), repository.getLfsDir());
        assertEquals(repositoryDir, repository.getPath());
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "templates"), repository.getTemplateDir());
    }

    @Test
    public void testAttachStorageToDeletedRepository() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));
        MoreFiles.write(repositoryDir.resolveSibling("1.deleted"), HIERARCHY_ID);

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertSame(StorageStatus.DELETED, shard.attachStorage(repository));
    }

    @Test
    public void testAttachStorageToDeletedRepositoryLong() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));
        MoreFiles.write(repositoryDir.resolveSibling("1.deleted"), HIERARCHY_ID);

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertSame(StorageStatus.DELETED, shard.attachStorage(repository));
    }

    @Test
    public void testAttachStorageToNonexistentRepository() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertSame(StorageStatus.NONEXISTENT, shard.attachStorage(repository));

        //This doesn't verify all the paths like testAttachStorage, but we do check path is set
        assertEquals(shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID).resolve("1"), repository.getPath());
    }

    @Test
    public void testAttachStorageToNonexistentRepositoryLong() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertSame(StorageStatus.NONEXISTENT, shard.attachStorage(repository));

        //This doesn't verify all the paths like testAttachStorage, but we do check path is set
        assertEquals(shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID).resolve("1"), repository.getPath());
    }

    @Test
    public void testAttachStorageToOccludedRepository() throws IOException {
        Path repositoriesDir = MoreFiles.mkdir(shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID));
        Path repositoryDir = repositoriesDir.resolve("1");
        MoreFiles.touch(repositoryDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertSame(StorageStatus.OCCLUDED, shard.attachStorage(repository));
    }

    @Test
    public void testAttachStorageToOccludedRepositoryLong() throws IOException {
        Path repositoriesDir = MoreFiles.mkdir(shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID));
        Path repositoryDir = repositoriesDir.resolve("1");
        MoreFiles.touch(repositoryDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertSame(StorageStatus.OCCLUDED, shard.attachStorage(repository));
    }

    @Test
    public void testAttachStorageWhenShardNotAvailable() {
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID, HIERARCHY_ID, "1");

        assertSame(shard.getState(), STARTING);
        DataStoreUnavailableException e = assertThrows(DataStoreUnavailableException.class,
                () -> shard.attachStorage(repository));

        assertThat(e.getMessage(), containsString("Data store is not available"));
    }

    @Test
    public void testAttachStorageWhenShardNotAvailableLong() {
        DefaultManagedRepository repository = new DefaultManagedRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");

        assertSame(shard.getState(), STARTING);
        DataStoreUnavailableException e = assertThrows(DataStoreUnavailableException.class,
                () -> shard.attachStorage(repository));

        assertThat(e.getMessage(), containsString("Data store is not available"));
    }

    @Test
    public void testClearMarkPartitionDeletedRemovesPartitionFromStreaming() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        Partition partition = new Partition(PARTITION_ID);
        shard.markPartitionDeleted(partition);

        shard.clearPartitionDeleted(partition);

        Map<Partition, Instant> deletedPartitions = new HashMap<>();
        shard.streamDeletedPartitions(deletedPartitions::put);
        assertNull(shard.getPartitionDeletedTime(partition));

        assertEquals(0, deletedPartitions.size());
    }

    @Test
    public void testDeletePartitionFilesAlsoDeletesMarkerFile() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir, "1");

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        Partition partition = new Partition(PARTITION_ID);
        shard.markPartitionDeleted(partition);
        ManagedShard.DeleteTask task = shard.deletePartitionFiles(partition);

        assertTrue(task.execute());

        assertFalse(Files.isDirectory(repositoriesDir));
        try (Stream<Path> files = Files.list(shard.getRootDir().resolve("data"))) {
            List<Path> filesList = files.collect(Collectors.toList());
            assertThat(filesList, emptyIterable());
        }
    }

    @Test
    public void testDeletePartitionFilesDeletesAllPartitionFiles() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        Path repository1Dir = MoreFiles.mkdir(repositoriesDir, "1");
        Path repository2Dir = MoreFiles.mkdir(repositoriesDir, "2");
        Path repository3Dir = MoreFiles.mkdir(repositoriesDir, "3");
        MoreFiles.mkdir(repository1Dir, "subdirectory");
        MoreFiles.mkdir(repository2Dir, "subdirectory");
        MoreFiles.mkdir(repository3Dir, "subdirectory");
        Files.createFile(repository1Dir.resolve("file"));
        Files.createFile(repository2Dir.resolve("file"));
        Files.createFile(repository3Dir.resolve("file"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        ManagedShard.DeleteTask task = shard.deletePartitionFiles(new Partition(PARTITION_ID));

        task.execute();

        assertFalse(Files.isDirectory(repositoriesDir));
        try (Stream<Path> files = Files.list(shard.getRootDir().resolve("data"))) {
            List<Path> filesList = files.collect(Collectors.toList());
            assertThat(filesList, emptyIterable());
        }
    }

    @Test
    public void testGetHooksDir() {
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "hooks"), shard.getHooksDir());
    }

    @Test
    public void testGetLfsDir() {
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID, HIERARCHY_ID, "lfs"),
                shard.getLfsDir(PARTITION_ID, HIERARCHY_ID));
    }

    @Test
    public void testGetLfsDirLong() {
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID_LONG, HIERARCHY_ID, "lfs"),
                shard.getLfsDir(PARTITION_ID_LONG, HIERARCHY_ID));
    }

    @Test
    public void testGetRepositoriesDir() {
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID, HIERARCHY_ID, "repos"),
                shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID));
    }

    @Test
    public void testGetRepositoriesDirLong() {
        assertEquals(MoreFiles.resolve(rootDir, "data", PARTITION_ID_LONG, HIERARCHY_ID, "repos"),
                shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID));
    }

    @Test
    public void testGetRepository() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        ManagedRepository repository = shard.getRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertNotNull(repository);
        //This doesn't verify all the paths like testAttachStorage, but we do check path is set
        assertEquals(repositoryDir, repository.getPath());
    }

    @Test
    public void testGetRepositoryLong() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID);
        Path repositoryDir = MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));
        ManagedRepository repository = shard.getRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertNotNull(repository);
        //This doesn't verify all the paths like testAttachStorage, but we do check path is set
        assertEquals(repositoryDir, repository.getPath());
    }

    @Test
    public void testGetRepositoryUnattached() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));
        ManagedRepository repository = shard.getRepository(PARTITION_ID, HIERARCHY_ID, "1");
        assertNull(repository);
    }

    @Test
    public void testGetRepositoryUnattachedLong() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));
        ManagedRepository repository = shard.getRepository(PARTITION_ID_LONG, HIERARCHY_ID, "1");
        assertNull(repository);
    }

    @Test
    public void testGetRootDir() {
        assertEquals(rootDir, shard.getRootDir());
    }

    @Test
    public void testGetSharedObjects() {
        SharedObjectsRepository repository = shard.getSharedObjects(PARTITION_ID, HIERARCHY_ID);
        assertNotNull(repository);
        assertEquals(shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID).resolve(SHARED_OBJECTS),
                repository.getPath());
    }

    @Test
    public void testGetSharedObjectsLong() {
        SharedObjectsRepository repository = shard.getSharedObjects(PARTITION_ID_LONG, HIERARCHY_ID);
        assertNotNull(repository);
        assertEquals(shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID).resolve(SHARED_OBJECTS),
                repository.getPath());
    }

    @Test
    public void testGetState() {
        //ManagedShards initialize in STARTING state until START is called
        assertSame(ShardState.STARTING, shard.getState());
    }

    @Test
    public void testGetTemplateDir() {
        assertEquals(MoreFiles.resolve(rootDir, "config", "default", "templates"), shard.getTemplateDir());
    }

    @Test
    public void testGetUsableSpace() throws IOException {
        //This verifies that we get back a positive size and no exception is thrown. Trying to validate
        //the number any more closely just ends up flaky
        assertThat(shard.getUsableSpace(), greaterThan(0L));
    }

    @Test
    public void testMarkPartitionDeletedStreamsDeletedPartition() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        Partition partition = new Partition(PARTITION_ID);
        Instant now = Instant.ofEpochSecond(Instant.now().getEpochSecond());
        shard.markPartitionDeleted(partition);

        Map<Partition, Instant> deletedPartitions = new HashMap<>();
        shard.streamDeletedPartitions(deletedPartitions::put);

        Instant partitionDeletedTime = shard.getPartitionDeletedTime(partition);

        assertNotNull(partitionDeletedTime);
        assertThat(partitionDeletedTime.getEpochSecond(), greaterThanOrEqualTo(now.getEpochSecond()));

        assertEquals(1, deletedPartitions.size());
        Map.Entry<Partition, Instant> entry = deletedPartitions.entrySet().iterator().next();
        assertEquals(partition, entry.getKey());
        assertThat(entry.getValue().getEpochSecond(), greaterThanOrEqualTo(now.getEpochSecond()));
    }

    @Test
    public void testMarkPartitionDeletedTwiceDoesntThrow() {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve("1"));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        Partition partition = new Partition(PARTITION_ID);
        shard.markPartitionDeleted(partition);
        shard.markPartitionDeleted(partition);
    }

    @Test
    public void testStart() throws IOException {
        //Create a few partition directories
        Path dataDir = MoreFiles.mkdir(rootDir, "data");
        MoreFiles.mkdir(dataDir, "0000");
        MoreFiles.mkdir(dataDir, "0001");
        MoreFiles.mkdir(dataDir, "0010");
        MoreFiles.mkdir(dataDir, "fffffffffffffff");

        Path hooksDir = MoreFiles.mkdir(shard.getHooksDir());
        Path newHook = hooksDir.resolve("new");
        Path oldHook = hooksDir.resolve("old");

        //Create an "old" file in hooksDir to verify it's pruned
        MoreFiles.write(oldHook, "Some test content", StandardOpenOption.CREATE_NEW);
        //Modify the modification timestamp to ensure the file is old
        BasicFileAttributeView attributes = Files.getFileAttributeView(oldHook, BasicFileAttributeView.class);
        FileTime lastModified = attributes.readAttributes().lastModifiedTime();
        lastModified = FileTime.fromMillis(lastModified.toMillis() - TimeUnit.HOURS.toMillis(1L));
        attributes.setTimes(lastModified, lastModified, lastModified);

        GitHookService hookService = mock(GitHookService.class);
        doAnswer(invocation -> {
            //Write a "new" hook to verify it's not pruned
            MoreFiles.write(newHook, "More test content", StandardOpenOption.CREATE_NEW);

            return null;
        }).when(hookService).installHooks(eq(hooksDir));

        shard.start(BUILD_VERSION, hookService);

        try {
            //Verify the shard transitioned to AVAILABLE
            assertSame(AVAILABLE, shard.getState());
            //Verify the expected partitions were loaded
            assertEquals(ImmutableSet.of("0000", "0001", "0010", "fffffffffffffff"),
                    shard.getPartitions().stream()
                            .map(Partition::getHex)
                            .collect(ImmutableSet.toImmutableSet()));
            //Verify system-config was unpacked for both default and shared-objects repositories
            assertTrue("Default system-config was not unpacked",
                    Files.isRegularFile(MoreFiles.resolve(rootDir, "config", "default", "system-config")));
            assertTrue("Shared objects system-config was not unpacked",
                    Files.isRegularFile(MoreFiles.resolve(rootDir, "config", "shared-objects", "system-config")));
            //Verify the stale hook was pruned and the new one wasn't
            assertFalse("The stale hook was not pruned", Files.exists(oldHook));
            assertTrue("The new hook was pruned", Files.isRegularFile(newHook));
        } finally {
            shard.stop();

            assertSame(ShardState.STOPPED, shard.getState());
        }
    }

    @Test
    public void testStartWithMajorDowngrade() throws IOException {
        //Write properties indicating the shard is one major version higher than the current one
        new DataStoreProperties(BUILD_VERSION, new Version(LAYOUT_VERSION.getMajor() + 1), TYPE).store(rootDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        assertSame(ShardState.INCOMPATIBLE, shard.getState());
        logListener.assertWarn("requires restoring from a backup");
    }

    @Test
    public void testStartWithMajorUpgrade() throws IOException {
        //Write properties indicating the shard is one major version lower than the current one
        new DataStoreProperties(BUILD_VERSION, new Version(LAYOUT_VERSION.getMajor() - 1), TYPE).store(rootDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        assertSame(AVAILABLE, shard.getState());
        logListener.assertInfo("will not be possible to downgrade");
    }

    @Test
    public void testStartWithMinorDowngrade() throws IOException {
        //Write properties indicating the shard is one minor version ahead on the same major version
        new DataStoreProperties(BUILD_VERSION,
                new Version(LAYOUT_VERSION.getMajor(), LAYOUT_VERSION.getMinor() + 1),
                TYPE).store(rootDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        assertSame(AVAILABLE, shard.getState());
        logListener.assertWarn("will not roll back the home directory");
    }

    @Test
    public void testStartWithMinorUpgrade() throws IOException {
        //If the current minor version is 0, there's no lower minor version in the same major version
        //that we can test against, so the test needs to be skipped. This uses assumeFalse so the test
        //will show as skipped rather than passing
        assumeFalse("Minor upgrades can't be tested when the current layout has a 0 minor version",
                LAYOUT_VERSION.getMinor() == 0);

        //Write properties indicating the shard is one minor version behind on the same major version
        new DataStoreProperties(BUILD_VERSION,
                new Version(LAYOUT_VERSION.getMajor(), LAYOUT_VERSION.getMinor() - 1),
                TYPE).store(rootDir);

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        assertSame(AVAILABLE, shard.getState());
        logListener.assertInfo("will be available while the upgrade is performed");
    }

    @Test
    public void testStartWithNewerBuild() throws IOException {
        new DataStoreProperties(BUILD_VERSION, LAYOUT_VERSION, TYPE).store(rootDir);

        Version newBuild = new Version(BUILD_VERSION.getMajor() + 1);
        shard.start(newBuild, mock(GitHookService.class));

        assertSame(AVAILABLE, shard.getState());

        DataStoreProperties properties = DataStoreProperties.load(rootDir);
        assertNotNull(properties);
        assertEquals(newBuild, properties.getBuild());
    }

    @Test
    public void testStartWithWrongType() throws IOException {
        //Write properties with an unexpected type
        new DataStoreProperties(BUILD_VERSION, LAYOUT_VERSION, "legacy").store(rootDir);

        IllegalStateException e = assertThrows(IllegalStateException.class,
                () -> shard.start(BUILD_VERSION, mock(GitHookService.class)));
        assertThat(e.getMessage(), containsString("is a legacy shard, not a"));
    }

    @Test
    public void testStreamRepositories() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve("1"));
        MoreFiles.mkdir(repositoriesDir.resolve("2"));
        MoreFiles.mkdir(repositoriesDir.resolve("3"));
        //Create a file to verify files are skipped
        MoreFiles.touch(repositoriesDir.resolve("file"));
        //Create a deleted repository to verify those are skipped as well
        Path deletedRepoDir = MoreFiles.mkdir(repositoriesDir.resolve("99"));
        MoreFiles.write(deletedRepoDir.resolveSibling("99.deleted"), HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve(SHARED_OBJECTS));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        ImmutableSet.Builder<String> builder = ImmutableSet.builder();
        shard.streamRepositories(PARTITION_ID, HIERARCHY_ID, repository -> builder.add(repository.getId()));

        assertEquals(ImmutableSet.of(
                ManagedRepository.formatId(PARTITION_ID, HIERARCHY_ID, "1"),
                ManagedRepository.formatId(PARTITION_ID, HIERARCHY_ID, "2"),
                ManagedRepository.formatId(PARTITION_ID, HIERARCHY_ID, "3")), builder.build());
    }

    @Test
    public void testStreamRepositoriesLong() throws IOException {
        Path repositoriesDir = shard.getRepositoriesDir(PARTITION_ID_LONG, HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve("1"));
        MoreFiles.mkdir(repositoriesDir.resolve("2"));
        MoreFiles.mkdir(repositoriesDir.resolve("3"));
        //Create a file to verify files are skipped
        MoreFiles.touch(repositoriesDir.resolve("file"));
        //Create a deleted repository to verify those are skipped as well
        Path deletedRepoDir = MoreFiles.mkdir(repositoriesDir.resolve("99"));
        MoreFiles.write(deletedRepoDir.resolveSibling("99.deleted"), HIERARCHY_ID);
        MoreFiles.mkdir(repositoriesDir.resolve(SHARED_OBJECTS));

        shard.start(BUILD_VERSION, mock(GitHookService.class));

        ImmutableSet.Builder<String> builder = ImmutableSet.builder();
        shard.streamRepositories(PARTITION_ID_LONG, HIERARCHY_ID, repository -> builder.add(repository.getId()));

        assertEquals(ImmutableSet.of(
                ManagedRepository.formatId(PARTITION_ID_LONG, HIERARCHY_ID, "1"),
                ManagedRepository.formatId(PARTITION_ID_LONG, HIERARCHY_ID, "2"),
                ManagedRepository.formatId(PARTITION_ID_LONG, HIERARCHY_ID, "3")), builder.build());
    }

    @Test
    public void testStreamRepositoriesNonExistentPartition() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));

        @SuppressWarnings("unchecked")
        Consumer<Repository> callback = mock(Consumer.class);
        shard.streamRepositories("fffffffffffffff", "11111111113333333333", callback);

        verifyNoInteractions(callback);
    }

    @Test
    public void testStreamRepositoriesNonexistentHierarchy() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));

        @SuppressWarnings("unchecked")
        Consumer<Repository> callback = mock(Consumer.class);
        shard.streamRepositories(PARTITION_ID, "11111111112222222222", callback);

        verifyNoInteractions(callback);
    }

    @Test
    public void testStreamRepositoriesNonexistentHierarchyLong() {
        shard.start(BUILD_VERSION, mock(GitHookService.class));

        @SuppressWarnings("unchecked")
        Consumer<Repository> callback = mock(Consumer.class);
        shard.streamRepositories(PARTITION_ID_LONG, "11111111112222222222", callback);

        verifyNoInteractions(callback);
    }
}
