package com.atlassian.bitbucket.mesh.management;

import com.atlassian.bitbucket.mesh.AbstractGrpcServiceTest;
import com.atlassian.bitbucket.mesh.MockResponseObserver;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.layout.StorageLayout;
import com.atlassian.bitbucket.mesh.network.*;
import com.atlassian.bitbucket.mesh.repository.ManagedRepository;
import com.atlassian.bitbucket.mesh.repository.ReplicaState;
import com.atlassian.bitbucket.mesh.repository.ReplicaStateRegistry;
import com.atlassian.bitbucket.mesh.repository.RepositoryAttributeManager;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.test.annotation.DirtiesRepo;
import com.atlassian.bitbucket.mesh.test.annotation.Repo;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Collections2;
import com.google.common.collect.ImmutableMap;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import org.apache.commons.compress.archivers.ArchiveEntry;
import org.apache.commons.compress.archivers.ArchiveInputStream;
import org.apache.commons.compress.archivers.zip.ZipArchiveInputStream;
import org.apache.commons.io.IOUtils;
import org.junit.After;
import org.junit.Test;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.SequenceInputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.layout.StorageLayout.MESH_PROPERTIES;
import static com.atlassian.bitbucket.mesh.management.SupportUtils.APPLICATION_PROPERTIES_JSON;
import static com.atlassian.bitbucket.mesh.rpc.ManagedConfigurationProperties.*;
import static java.util.stream.Collectors.toList;
import static org.hamcrest.Matchers.*;

@SuppressWarnings("ResultOfMethodCallIgnored")
public class GrpcManagementServiceTest extends AbstractGrpcServiceTest {

    private static final String UNKNOWN_REPOSITORY = "p/0000/h/abc123abc123abc123ab/r/0001112223";

    @Resource
    private ApplicationConfig applicationConfig;
    @Resource
    private RepositoryAttributeManager attributeManager;
    @Resource
    private MeshManager meshManager;
    @Resource
    private NodeManager nodeManager;
    @Resource
    private ReplicaStateRegistry replicaStateRegistry;
    @Repo
    private ManagedRepository repository;
    @Resource
    private SidebandManager sidebandManager;
    @Resource
    private StorageLayout storageLayout;

    @After
    public void tearDown() {
        attributeManager.replaceReadOnly(Collections.emptySet());
    }

    @Test(timeout = 2000L)
    public void testDrainAll() throws TimeoutException, InterruptedException {
        //Start a long running repository operation so we can test draining
        MockResponseObserver<RpcStreamFragment> sourceObserver = new MockResponseObserver<>();
        StreamObserver<RpcCatFileStdinFragment> observer = plumbingServiceStub().catFileBatchCheck(sourceObserver);
        //Access a repository to increase the active request count for its partition
        observer.onNext(RpcCatFileStdinFragment.newBuilder()
                .setRequest(RpcCatFileBatchRequest.newBuilder()
                        .setRepository(repository.getId()))
                .build());
        //Wait for the RPC to respond with a message indicating that the process is ready
        sourceObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);

        //Now request the node to be drained
        MockResponseObserver<RpcDrainResponseFragment> drainObserver = new MockResponseObserver<>();
        stub().drain(RpcDrainRequest.newBuilder()
                .setAllRequest(RpcAllRequest.getDefaultInstance())
                .build(), drainObserver);

        //The node should not drain because the repair operation is still running
        assertThrows(TimeoutException.class, () -> drainObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS));

        //Finish the long running RPC and wait for the repository hierarchy to return as drained
        observer.onCompleted();

        RpcDrainResponseFragment fragment = drainObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);
        assertTrue(fragment.hasDrainedResource());
        assertEquals(RpcDrainedResource.getDefaultInstance(), fragment.getDrainedResource());

        drainObserver.await();
    }

    @Test(timeout = 2000L)
    public void testDrainHierarchy() throws TimeoutException, InterruptedException {
        //Start a long running repository operation so we can test draining
        MockResponseObserver<RpcStreamFragment> sourceObserver = new MockResponseObserver<>();
        StreamObserver<RpcCatFileStdinFragment> observer = plumbingServiceStub().catFileBatchCheck(sourceObserver);
        //Access a repository to increase the active request count for its partition
        observer.onNext(RpcCatFileStdinFragment.newBuilder()
                .setRequest(RpcCatFileBatchRequest.newBuilder()
                        .setRepository(repository.getId()))
                .build());
        //Wait for the RPC to respond with a message indicating that the process is ready
        sourceObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);

        //Now request the hierarchy to be drained
        MockResponseObserver<RpcDrainResponseFragment> drainObserver = new MockResponseObserver<>();
        stub().drain(RpcDrainRequest.newBuilder()
                .setHierarchiesRequest(RpcHierarchiesRequest.newBuilder()
                        .addHierarchies(RpcHierarchyRequest.newBuilder()
                                .setHierarchy(repository.getHierarchyId())
                                .setPartition(repository.getPartition()))
                        .addHierarchies(RpcHierarchyRequest.newBuilder()
                                .setHierarchy("aaaaaaaaaabbbbbbbbbb") // hierarchy with no active requests
                                .setPartition(repository.getPartition())))
                .build(), drainObserver);

        //The non-existent hierarchy should drain immediately
        RpcDrainResponseFragment fragment = drainObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);
        assertTrue(fragment.hasDrainedResource());
        RpcDrainedResource resource = fragment.getDrainedResource();
        assertEquals("aaaaaaaaaabbbbbbbbbb", resource.getHierarchy());
        assertEquals(repository.getPartition(), resource.getPartition());

        //The repository hierarchy should not drain because the repair operation is still running
        assertThrows(TimeoutException.class, () -> drainObserver.awaitFragment(1, 150, TimeUnit.MILLISECONDS));

        //Finish the long-running RPC and wait for the repository hierarchy to return as drained
        observer.onCompleted();

        fragment = drainObserver.awaitFragment(1, 150, TimeUnit.MILLISECONDS);
        assertTrue(fragment.hasDrainedResource());
        resource = fragment.getDrainedResource();
        assertEquals(repository.getHierarchyId(), resource.getHierarchy());
        assertEquals(repository.getPartition(), resource.getPartition());

        drainObserver.await();
    }

    @Test(timeout = 20000L)
    public void testDrainPartitions() throws TimeoutException, InterruptedException {
        //Start a long running repository operation so we can test draining
        MockResponseObserver<RpcStreamFragment> sourceObserver = new MockResponseObserver<>();
        StreamObserver<RpcCatFileStdinFragment> observer = plumbingServiceStub().catFileBatchCheck(sourceObserver);
        //Access a repository to increase the active request count for its partition
        observer.onNext(RpcCatFileStdinFragment.newBuilder()
                .setRequest(RpcCatFileBatchRequest.newBuilder()
                        .setRepository(repository.getId()))
                .build());
        //Wait for the RPC to respond with a message indicating that the process is ready
        sourceObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);

        //Now request the partitions to be drained
        MockResponseObserver<RpcDrainResponseFragment> drainObserver = new MockResponseObserver<>();
        String otherPartition = "0000".equals(repository.getPartition()) ? "0001" : "0000";
        stub().drain(RpcDrainRequest.newBuilder()
                .setPartitionsRequest(RpcPartitionsRequest.newBuilder()
                        .addPartitions(repository.getPartition())
                        .addPartitions(otherPartition))
                .build(), drainObserver);

        //The other partition should drain immediately
        RpcDrainResponseFragment fragment = drainObserver.awaitFragment(0, 150, TimeUnit.MILLISECONDS);
        assertTrue(fragment.hasDrainedResource());
        RpcDrainedResource resource = fragment.getDrainedResource();
        assertFalse(resource.hasHierarchy());
        assertEquals(otherPartition, resource.getPartition());

        //The repository partition should not drain because the repair operation is still running
        assertThrows(TimeoutException.class, () -> drainObserver.awaitFragment(1, 150, TimeUnit.MILLISECONDS));

        //Finish the repair RPC and wait for the repository partition to return as drained
        observer.onCompleted();

        fragment = drainObserver.awaitFragment(1, 150, TimeUnit.MILLISECONDS);
        assertTrue(fragment.hasDrainedResource());
        resource = fragment.getDrainedResource();
        assertFalse(resource.hasHierarchy());
        assertEquals(repository.getPartition(), resource.getPartition());

        drainObserver.await();
    }

    @Test
    public void testGetConfiguration() {
        long tempExpiryHours = applicationConfig.getStorageConfig().getTempExpiry().toHours();
        try {
            RpcNode extraNode = RpcNode.newBuilder()
                    .setId("923")
                    .setName("extra-node")
                    .setRpcUrl("https://extra.mesh.example.com:7777")
                    .build();
            String updatedExpiry = Long.toString(tempExpiryHours + 1); //non-default value
            //Update the configuration to make sure there's at least one inconsistent replica and a custom
            //configuration property that is returned to us.
            RpcPartitionReplicaMigration migration = RpcPartitionReplicaMigration.newBuilder()
                    .setPartition("0000")
                    .setSourceNodeId(meshManager.getLocalNode().getId())
                    .setTargetNodeId("923")
                    .build();
            RpcSetConfigurationRequest request = RpcSetConfigurationRequest.newBuilder()
                    .setNodeConfiguration(RpcNodeConfiguration.newBuilder()
                            .putProperties("storage.temp-expiry", updatedExpiry))
                    //set the current topology, but add a partition migration
                    .setTopology(RpcTopology.newBuilder()
                            .addAllNodes(Collections2.transform(meshManager.getNodes(), MeshNode::toNode))
                            .addAllPartitions(
                                    Collections2.transform(meshManager.getPartitions(), Partition::toPartition))
                            .addAllPartitionMigrations(Collections2.transform(meshManager.getPartitionMigrations(),
                                    PartitionReplicaMigration::toPartitionReplicaMigration))
                            .addNodes(extraNode)
                            .addPartitionMigrations(migration))
                    .addInconsistentReplicas(RpcInconsistentReplica.newBuilder()
                            .setRepository(repository.getId())
                            .setVersion(repository.getVersion() + 1L))
                    .build();
            blockingStub().setConfiguration(request);

            RpcGetConfigurationResponse response = blockingStub()
                    .getConfiguration(RpcGetConfigurationRequest.getDefaultInstance());

            assertNotNull(response);
            //Verify that the custom property is returned
            Map<String, String> configProperties = response.getNodeConfiguration().getPropertiesMap();
            assertEquals(updatedExpiry, configProperties.get("storage.temp-expiry"));
            //Verify that a valid topology is returned
            RpcTopology topology = response.getTopology();
            //Verify nodes list
            assertThat(topology.getNodesCount(), equalTo(request.getTopology().getNodesCount()));
            for (RpcNode node : topology.getNodesList()) {
                assertThat(node.getId(), not(emptyString()));
                assertThat(node.getName(), not(emptyString()));
                assertThat(node.getRpcUrl(), not(emptyString()));
            }
            //Verify partitions list
            assertThat(topology.getPartitionsCount(), equalTo(request.getTopology().getPartitionsCount()));
            for (RpcPartition partition : topology.getPartitionsList()) {
                assertThat(partition.getId(), not(emptyString()));
                assertThat(partition.getReplicasCount(), greaterThan(0));
                RpcPartitionReplica partitionReplica = partition.getReplicas(0);
                assertThat(partitionReplica.getId(), not(emptyString()));
                assertThat(partitionReplica.getNodeId(), not(emptyString()));
            }
            //Verify the inconsistent replica list (there may be more inconsistent replicas thanks to other tests)
            assertThat(response.getInconsistentReplicasCount(),
                    greaterThanOrEqualTo(request.getInconsistentReplicasCount()));
            RpcInconsistentReplica inconsistentReplica = response.getInconsistentReplicasList().stream()
                    .filter(item -> item.getRepository().equals(repository.getId()))
                    .findFirst()
                    .orElse(null);
            assertNotNull(inconsistentReplica);
            assertEquals(repository.getVersion() + 1, inconsistentReplica.getVersion());
            //Verify the migrations list
            assertThat(topology.getPartitionMigrationsCount(),
                    equalTo(request.getTopology().getPartitionMigrationsCount()));
            assertTrue(topology.getPartitionMigrationsList().contains(migration));
        } finally {
            //Mark the repository as consistent again
            replicaStateRegistry.setState(repository.getId(), ReplicaState.CONSISTENT, repository.getVersion(), true);
            //Restore the original topology, discarding the migration
            meshManager.updateTopology(new UpdateTopologyParameters.Builder()
                    .nodes(Collections2.filter(meshManager.getNodes(), node -> !"923".equals(node.getId())))
                    .partitions(meshManager.getPartitions())
                    .build());
        }
    }

    @Test
    public void testGetDataStores() {
        RpcGetDataStoresResponse dataStores = blockingStub().getDataStores(RpcGetDataStoresRequest.getDefaultInstance());
        assertEquals(3, dataStores.getDataStoresCount());
        Set<RpcDataStore.Type> types = EnumSet.allOf(RpcDataStore.Type.class);
        for (RpcDataStore store : dataStores.getDataStoresList()) {
            assertTrue("Should have at most one of each type", types.remove(store.getType()));
            assertTrue(store.getType() + " should have at least one shard", store.getShardsCount() > 0);
        }
    }

    @Test
    public void testGetSupportZip() throws Exception {
        doTestGetSupportZip(true);
    }

    @Test
    public void testGetSupportZipWithNoConfig() throws Exception {
        doTestGetSupportZip(false);
    }

    @Test
    public void testHeartbeat() {
        RpcHeartbeat heartbeat = blockingStub().heartbeat(RpcHeartbeat.getDefaultInstance());
        assertNotNull(heartbeat);
    }

    @Test
    public void testMarkInconsistent() {
        String repositoryId = ManagedRepository.formatId(1, "abcdefabcdefabcdef00", "1231");
        blockingStub().markRepositoriesInconsistent(RpcMarkRepositoriesInconsistentRequest.newBuilder()
                .addInconsistentReplicas(RpcInconsistentReplica.newBuilder()
                        .setRepository(repositoryId)
                        .setVersion(1L))
                .build());

        assertEquals(ReplicaState.INCONSISTENT, replicaStateRegistry.getState(repositoryId));
    }

    @Test
    public void testReplaceReadOnlyRepositories() throws Exception {
        MockResponseObserver<RpcReplaceReadOnlyRepositoriesResponse> responseObserver = new MockResponseObserver<>();
        StreamObserver<RpcReplaceReadOnlyRepositoriesFragment> requestObserver =
                stub().replaceReadOnlyRepositories(responseObserver);

        try {
            requestObserver.onNext(RpcReplaceReadOnlyRepositoriesFragment.newBuilder()
                    .addRepositories(repository.getId())
                    .build());
            requestObserver.onNext(RpcReplaceReadOnlyRepositoriesFragment.newBuilder()
                    .addRepositories(UNKNOWN_REPOSITORY) //Will be ignored without causing failure
                    .build());
            requestObserver.onCompleted();

            responseObserver.await();
            assertEquals(1, responseObserver.getFragments().size());
            assertTrue(repository.isReadOnly());
        } finally {
            //Re-enable read/write mode
            repository.setReadOnly(false);
        }
    }

    @Test
    public void testReplaceReadOnlyRepositoriesOnError() {
        MockResponseObserver<RpcReplaceReadOnlyRepositoriesResponse> responseObserver = new MockResponseObserver<>();
        StreamObserver<RpcReplaceReadOnlyRepositoriesFragment> requestObserver =
                stub().replaceReadOnlyRepositories(responseObserver);

        requestObserver.onNext(RpcReplaceReadOnlyRepositoriesFragment.newBuilder()
                .addRepositories(repository.getId())
                .build());
        //Any call to onError instead of onComplete means the repositories sent in any of the fragments will not be
        //updated, the changes are only applied in onComplete
        requestObserver.onError(Status.CANCELLED.asRuntimeException());

        assertThrows(StatusRuntimeException.class, responseObserver::await);
        assertFalse(repository.isReadOnly());
    }

    @DirtiesRepo
    @Test
    public void testResetHierarchy() {
        RpcResetHierarchyResponse response = blockingStub().resetHierarchy(RpcResetHierarchyRequest.newBuilder()
                .setPartition(repository.getPartition())
                .setHierarchy(repository.getHierarchyId())
                .build());

        assertEquals(1, response.getResetCount());
    }

    @DirtiesRepo
    @Test
    public void testResetHierarchyThatDoesNotExist() {
        RpcResetHierarchyResponse response = blockingStub().resetHierarchy(RpcResetHierarchyRequest.newBuilder()
                .setPartition(repository.getPartition())
                .setHierarchy("00000000003333333333")
                .build());

        assertEquals(0, response.getResetCount());
    }

    @DirtiesRepo
    @Test
    public void testResetHierarchyOnPartitionThatDoesNotExist() {
        RpcResetHierarchyResponse response = blockingStub().resetHierarchy(RpcResetHierarchyRequest.newBuilder()
                .setPartition("fffff")
                .setHierarchy(repository.getHierarchyId())
                .build());

        assertEquals(0, response.getResetCount());
    }

    @Test
    public void testSetConfigurationEmpty() {
        blockingStub().setConfiguration(RpcSetConfigurationRequest.getDefaultInstance());
    }

    @Test
    public void testSetConfigurationPropertiesAndTopology() {
        doWithConfiguration(ignored -> {
            MeshNode localNode = meshManager.getLocalNode();
            MeshNode node1 = createNode(1);
            MeshNode node2 = createNode(2);

            RpcSetConfigurationResponse response =
                    blockingStub().setConfiguration(RpcSetConfigurationRequest.newBuilder()
                            .setNodeConfiguration(RpcNodeConfiguration.newBuilder()
                                    .putProperties(NODE_ID, localNode.getId())
                                    .putProperties(NODE_NAME, localNode.getName())
                                    .putProperties(NODE_RPC_URL, localNode.getRpcUrl()))
                            .setTopology(RpcTopology.newBuilder()
                                    .addNodes(node1.toNode())
                                    .addNodes(node2.toNode())
                                    .addNodes(localNode.toNode())
                                    .addPartitions(RpcPartition.newBuilder()
                                            .setId("0007")
                                            .addReplicas(RpcPartitionReplica.newBuilder()
                                                    .setId("replica-7-1")
                                                    .setNodeId(node2.getId()))
                                            .addReplicas(RpcPartitionReplica.newBuilder()
                                                    .setId("replica-7-2")
                                                    .setNodeId(node1.getId()))
                                            .addReplicas(RpcPartitionReplica.newBuilder()
                                                    .setId("replica-7-3")
                                                    .setNodeId(localNode.getId())))
                                    .addPartitionMigrations(RpcPartitionReplicaMigration.newBuilder()
                                            .setPartition("0007")
                                            .setSourceNodeId(node1.getId())
                                            .setTargetNodeId(node2.getId()))
                            )
                            .build());
            assertNotNull(response);

            assertThat(meshManager.getNodes(), containsInAnyOrder(node1, node2, localNode));
            assertEquals(1, meshManager.getPartitions().size());

            assertPartition("0007", ImmutableMap.of(
                            "replica-7-1", node2,
                            "replica-7-2", node1,
                            "replica-7-3", localNode),
                    meshManager.getPartition("0007"));
            assertThat(meshManager.getPartitionMigrations().size(), equalTo(1));
            PartitionReplicaMigration migration = meshManager.getPartitionMigrations().iterator().next();
            assertEquals("0007", migration.getPartitionId());
            assertEquals(node1, migration.getSourceNode());
            assertEquals(node2, migration.getTargetNode());
        });
    }

    @Test
    public void testSetConfigurationTopologyOnly() {
        doWithConfiguration(localNode -> {
            MeshNode node1 = createNode(1);
            MeshNode node2 = createNode(2);
            RpcTopology topology = RpcTopology.newBuilder()
                    .addNodes(localNode.toNode())
                    .addNodes(node1.toNode())
                    .addNodes(node2.toNode())
                    .addPartitions(RpcPartition.newBuilder()
                            .setId("0001")
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-1-1")
                                    .setNodeId(localNode.getId()))
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-1-2")
                                    .setNodeId(node1.getId()))
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-1-3")
                                    .setNodeId(node2.getId())))
                    .addPartitions(RpcPartition.newBuilder()
                            .setId("0003")
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-3-1")
                                    .setNodeId(node1.getId()))
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-3-2")
                                    .setNodeId(node2.getId()))
                            .addReplicas(RpcPartitionReplica.newBuilder()
                                    .setId("replica-3-3")
                                    .setNodeId(localNode.getId())))
                    .build();

            RpcSetConfigurationResponse response = blockingStub().setConfiguration(RpcSetConfigurationRequest.newBuilder()
                    .setTopology(topology)
                    .build());
            assertNotNull(response);

            assertThat(meshManager.getNodes(), containsInAnyOrder(localNode, node1, node2));
            assertEquals(2, meshManager.getPartitions().size());

            assertPartition("0001", ImmutableMap.of(
                            "replica-1-1", localNode,
                            "replica-1-2", node1,
                            "replica-1-3", node2),
                    meshManager.getPartition("0001"));
            assertPartition("0003", ImmutableMap.of(
                            "replica-3-1", node1,
                            "replica-3-2", node2,
                            "replica-3-3", localNode),
                    meshManager.getPartition("0003"));
        });
    }

    @Test
    public void testSetReplicaState() {
        try {
            blockingStub().setReplicaState(RpcSetReplicaStateRequest.newBuilder()
                    .setRepository(repository.getId())
                    .setReplicaState(RpcReplicaState.REPLICA_STATE_INCONSISTENT)
                    .setVersion(repository.getVersion() + 1L)
                    .build());

            assertEquals(ReplicaState.INCONSISTENT, replicaStateRegistry.getState(repository));

            //Marking it consistent with an older version should be ignored
            blockingStub().setReplicaState(RpcSetReplicaStateRequest.newBuilder()
                    .setRepository(repository.getId())
                    .setReplicaState(RpcReplicaState.REPLICA_STATE_CONSISTENT)
                    .setVersion(repository.getVersion())
                    .build());
            assertEquals(ReplicaState.INCONSISTENT, replicaStateRegistry.getState(repository));

            //Marking it consistent with an older version should be ignored, unless we set the force flag
            blockingStub().setReplicaState(RpcSetReplicaStateRequest.newBuilder()
                    .setForce(true)
                    .setRepository(repository.getId())
                    .setReplicaState(RpcReplicaState.REPLICA_STATE_CONSISTENT)
                    .setVersion(repository.getVersion())
                    .build());
            assertEquals(ReplicaState.CONSISTENT, replicaStateRegistry.getState(repository));

        } finally {
            //Restore the state to CONSISTENT
            replicaStateRegistry.setState(repository.getId(), ReplicaState.CONSISTENT, repository.getVersion(), true);
        }
    }

    @Test
    public void testSetRepositoryReadOnlyFailsWithUnknownRepository() {
        RpcSetRepositoryReadOnlyRequest request = RpcSetRepositoryReadOnlyRequest.newBuilder()
                .setRepository(UNKNOWN_REPOSITORY)
                .setReadOnly(true)
                .build();

        StatusRuntimeException thrown = assertThrows(StatusRuntimeException.class,
                () -> blockingStub().setRepositoryReadOnly(request));
        assertSame(Status.Code.NOT_FOUND, thrown.getStatus().getCode());
    }

    @Test
    public void testSetRepositoryReadOnlyFalse() {
        repository.setReadOnly(true);
        RpcSetRepositoryReadOnlyRequest request = RpcSetRepositoryReadOnlyRequest.newBuilder()
                .setRepository(repository.getId())
                .setReadOnly(false)
                .build();

        blockingStub().setRepositoryReadOnly(request);
        assertFalse(repository.isReadOnly());
    }

    @Test
    public void testSetRepositoryReadOnlyTrue() {
        RpcSetRepositoryReadOnlyRequest request = RpcSetRepositoryReadOnlyRequest.newBuilder()
                .setRepository(repository.getId())
                .setReadOnly(true)
                .build();

        blockingStub().setRepositoryReadOnly(request);
        assertTrue(repository.isReadOnly());
    }

    @Test
    public void testSideband() throws Exception {
        RpcSidebandFollowerFragment followerFragment = RpcSidebandFollowerFragment.newBuilder()
                .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                .build();
        //This fragment should be queued, so it will still be sent later
        assertTrue(sidebandManager.send(followerFragment));

        CountDownLatch firstFragment = new CountDownLatch(1);
        MockResponseObserver<RpcSidebandFollowerFragment> responseObserver =
                new MockResponseObserver<RpcSidebandFollowerFragment>() {

                    @Override
                    public void onNext(RpcSidebandFollowerFragment value) {
                        firstFragment.countDown();

                        super.onNext(value);
                    }
                };

        StreamObserver<RpcSidebandLeaderFragment> requestObserver = stub().sideband(responseObserver);
        //Wait for the queued heartbeat to be sent. We need to this to ensure the sideband has
        //actually been opened, since requests are async, before we try to send another
        assertTrue(firstFragment.await(10L, TimeUnit.SECONDS));

        //Send a heartbeat from the "control plane" to the client, which should be echoed back
        RpcSidebandLeaderFragment leaderFragment = RpcSidebandLeaderFragment.newBuilder()
                .setHeartbeat(RpcHeartbeat.getDefaultInstance())
                .build();
        requestObserver.onNext(leaderFragment);

        //Verify we can send a heartbeat to the "control plane" now
        try {
            assertTrue(sidebandManager.send(followerFragment));
        } finally {
            //Send another heartbeat from the "control plane" and then complete the RPC to simulate
            //the "control plane" disconnecting
            requestObserver.onNext(leaderFragment);
            requestObserver.onCompleted();
            responseObserver.await();

            //The "control plane" sent two heartbeats, so we should receive 2 back, plus two more for
            //the follower-initiated heartbeats (one before opening and one after)
            //But why greaterThanOrEqualTo? Because _other tests_ can end up queueing fragments, if
            //their logic tries to use the sideband, so we might get more fragments than just those
            //that are exchanged by this test
            assertThat(responseObserver.getFragments().size(), greaterThanOrEqualTo(4));
        }
    }

    private static void assertPartition(String expectedId, Map<String, MeshNode> expectedReplicas,
                                        Partition partition) {
        assertNotNull("Partition not found", partition);
        assertEquals("Partition ID does not match", expectedId, partition.getId());

        Map<String, MeshNode> actualReplicas = partition.getReplicas().stream()
                .collect(Collectors.toMap(PartitionReplica::getId, PartitionReplica::getNode));
        assertEquals("Partition replicas do not match", expectedReplicas, actualReplicas);
    }

    private static MeshNode createNode(int id) {
        return new MeshNode.Builder(Integer.toString(id))
                .name("node-" + id)
                .rpcUrl("http://node-" + id + ":7777")
                .build();
    }

    private static Map<String, String> parseSupportZipStream(List<RpcStreamChunk> chunks) throws Exception {
        ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
        try (SequenceInputStream sequenceInputStream = new SequenceInputStream(Collections.enumeration(chunks.stream()
                .map(RpcStreamChunk::getData)
                .filter(chunkData -> !chunkData.isEmpty())
                .map(ByteString::newInput)
                .collect(toList())));
             ArchiveInputStream inputStream = new ZipArchiveInputStream(sequenceInputStream)) {
            ArchiveEntry entry;
            while ((entry = inputStream.getNextEntry()) != null) {
                if (entry.isDirectory()) {
                    continue;
                }
                String path = entry.getName();
                String contents = IOUtils.toString(inputStream, StandardCharsets.UTF_8);
                if (path != null && contents != null) {
                    builder.put(path, contents);
                }
            }
        }
        return builder.build();
    }

    private ManagementServiceGrpc.ManagementServiceBlockingStub blockingStub() {
        return ManagementServiceGrpc.newBlockingStub(managedChannel);
    }

    private void doWithConfiguration(Consumer<MeshNode> test) {
        MeshNode localNode = meshManager.getLocalNode();
        Collection<Partition> partitions = meshManager.getPartitions();
        try {
            test.accept(localNode);
        } finally {
            //Restore the original local node configuration, in case it was updated, and clear any
            //topology modifications
            nodeManager.updateConfiguration(new UpdateConfigurationParameters.Builder()
                    .property(NODE_ID, localNode.getId())
                    .property(NODE_NAME, localNode.getName())
                    .property(NODE_RPC_URL, localNode.getRpcUrl())
                    .topology(new UpdateTopologyParameters.Builder()
                            .nodes(localNode)
                            .partitions(partitions)
                            .build())
                    .build());
        }
    }

    private void doTestGetSupportZip(boolean withConfig) throws Exception {
        //Make sure a mesh.properties exists so we can verify whether it is included
        Path applicationProperties = storageLayout.getHomeDir().resolve(MESH_PROPERTIES);
        MoreFiles.touch(applicationProperties);
        //Write a log file to disk
        Path logFile = writeRandomBytesToLogs();

        Map<String, String> supportZip = performGetSupportZip(withConfig);

        assertFalse(supportZip.isEmpty());
        //Verify that the log file is present
        assertTrue(supportZip.keySet().stream().anyMatch(logFile::endsWith));
        //Verify if that the config file is (NOT) present
        assertEquals(withConfig, supportZip.keySet().stream().anyMatch(path -> path.endsWith(MESH_PROPERTIES)));

        //Verify that the application-properties.json is always present
        String applicationJson = supportZip.get(APPLICATION_PROPERTIES_JSON);
        assertNotNull("application-properties.json not found", applicationJson);
        assertThat(applicationJson, containsString("build-version"));
        //Verify that application-properties.json is valid JSON
        Map<String, Object> values = new ObjectMapper().readValue(applicationJson, LinkedHashMap.class);

        assertNotNull(values);
        Set<String> categories = values.keySet();
        assertTrue("application section is missing", categories.contains("application"));
        assertTrue("node section is missing", categories.contains("node"));
        assertTrue("jvm section is missing", categories.contains("jvm"));
        assertTrue("git section is missing", categories.contains("git"));
        assertTrue("java runtime section is missing", categories.contains("java-runtime-environment"));
        assertTrue("environment variable section is missing", categories.contains("environment-variables"));
        assertTrue("operating system section is missing", categories.contains("operating-system"));
        assertTrue("data store section is missing", categories.contains("data-stores"));
        assertTrue("inconsistent replica section is missing", categories.contains("inconsistent-replicas"));
    }

    private Map<String, String> performGetSupportZip(boolean withConfig) throws Exception {
        RpcGetSupportZipRequest request = RpcGetSupportZipRequest.newBuilder()
                .setConfig(withConfig)
                .setLogs(true)
                .build();
        MockResponseObserver<RpcStreamChunk> responseObserver = new MockResponseObserver<>();

        stub().getSupportZip(request, responseObserver);
        responseObserver.await();

        List<RpcStreamChunk> fragments = responseObserver.getFragments();
        assertFalse(fragments.isEmpty());

        return parseSupportZipStream(fragments);
    }

    private PlumbingServiceGrpc.PlumbingServiceStub plumbingServiceStub() {
        return PlumbingServiceGrpc.newStub(managedChannel);
    }

    private ManagementServiceGrpc.ManagementServiceStub stub() {
        return ManagementServiceGrpc.newStub(managedChannel);
    }

    private Path writeRandomBytesToLogs() throws IOException {
        //Generate some random bytes to be populated in the log files to ensure multiple onNext calls are made
        Random rnd = new Random();
        byte[] bytes = new byte[15000];
        rnd.nextBytes(bytes);

        Path logDir = storageLayout.getLogDir();
        Path supportFile = Files.createTempFile(logDir, "mesh-support-1", ".log");
        Files.write(supportFile, bytes, StandardOpenOption.WRITE);
        return supportFile;
    }
}