package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.config.ReplicationConfig;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteToken;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenHandler;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.GrpcExecutionManager;
import com.atlassian.bitbucket.mesh.git.builder.GitProcess;
import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.git.builder.indexpack.GitIndexPackBuilder;
import com.atlassian.bitbucket.mesh.git.builder.indexpack.IndexPackStdioHandler;
import com.atlassian.bitbucket.mesh.git.exception.CommandFailedException;
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.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.PartitionReplicaMigration;
import com.atlassian.bitbucket.mesh.process.OsProcess;
import com.atlassian.bitbucket.mesh.process.ProcessSummary;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcPrepareTransactionResponse.Vote;
import com.atlassian.bitbucket.mesh.test.ExecutionContextRule;
import com.atlassian.bitbucket.mesh.transaction.*;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableSet;
import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.mutable.MutableBoolean;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Answers;
import org.mockito.ArgumentCaptor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Collections;
import java.util.concurrent.*;

import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class RemoteTransactionLeaderTest {

    private static final String CONTENT_HASH = "content-hash";
    private static final String LEADER_ID = "leader-id";
    private static final String LOCAL_ID = "local-id";
    private static final String PARTITION_ID = StringUtils.leftPad(Integer.toHexString(123), 4, '0');
    private static final String REPOSITORY_ID = ManagedRepository.formatId(123, "hierarchy-id", "repo-id");
    private static final long TIME_PACK = 1000L;
    private static final long TIME_PREPARE = 1003L;
    private static final long TIME_COMMIT = 1006L;
    private static final long TIME_ROLLBACK = 1006L;
    private static final long TIME_SUMMARY = 1021L;
    private static final long TIMEOUT_MS = 2000L;
    private static final String TRANSACTION_ID = Transaction.formatId(LEADER_ID, "97361238");
    private static final long VERSION = 77L;

    @Rule
    public final ExecutionContextRule executionContextRule = new ExecutionContextRule(new GrpcExecutionManager());
    @Rule
    public final LogListener logListener = new LogListener("com.atlassian.bitbucket.mesh");
    @Rule
    public final TemporaryDirectory temporaryDirectory = new TemporaryDirectory();

    @Mock(answer = Answers.RETURNS_DEEP_STUBS)
    private GitProcessBuilderFactory builderFactory;
    private CompletableFuture<TransactionState> completionState;
    @Mock
    private EventPublisher eventPublisher;
    private RemoteTransactionLeader leader;
    @Mock
    private MeshNode leaderNode;
    @Mock
    private MeshNode localNode;
    @Mock
    private PreparedAction localOperation;
    @Mock
    private LogicalClock logicalClock;
    @Mock
    private MeshManager meshManager;
    private ReplicaState publishedState;
    @Mock
    private ReplicaStateRegistry replicaStateRegistry;
    @Mock
    private ManagedRepository repository;
    @Mock
    private RepositoryAttributeManager repositoryAttributeManager;
    @Mock
    private RepositoryHashService repositoryHashService;
    @Mock
    private RepositoryManager repositoryManager;
    private RequestObserver requestObserver;
    private RpcTransaction rpcTransaction;
    @Mock
    private RpcTransactionExecutor transactionExecutor;
    private TransactionManager transactionManager;
    private boolean writeTokenAvailable;
    private WriteTokenHandler writeTokenHandler;
    @Mock
    private WriteToken.Promise writeTokenPromise;
    @Mock
    private WriteTokenRegistry writeTokenRegistry;

    @Before
    public void setUp() {
        completionState = new CompletableFuture<>();

        when(meshManager.getLocalNode()).thenReturn(localNode);
        when(meshManager.getNodeById(LEADER_ID)).thenReturn(leaderNode);
        when(meshManager.getPartitionMigrations(anyString())).thenReturn(Collections.emptySet());
        when(repository.getId()).thenReturn(REPOSITORY_ID);
        when(repository.getPath()).thenReturn(temporaryDirectory.getRoot());
        when(repository.getPartition()).thenReturn(PARTITION_ID);
        when(repository.accept(any())).thenAnswer(invocation -> {
            RepositoryVisitor<?> visitor = invocation.getArgument(0);
            return visitor.visit(repository);
        });
        when(repositoryManager.findById(REPOSITORY_ID)).thenReturn(repository);

        lenient().when(leaderNode.getId()).thenReturn(LEADER_ID);
        lenient().when(localNode.getId()).thenReturn(LOCAL_ID);
        lenient().when(localNode.isLocal()).thenReturn(true);
        lenient().when(meshManager.getNodeById(LOCAL_ID)).thenReturn(localNode);

        ReplicationConfig replicationConfig = mock(ReplicationConfig.class);
        lenient().when(replicationConfig.getQuorumMode()).thenReturn(QuorumMode.UNANIMOUS);
        ApplicationConfig config = mock(ApplicationConfig.class);
        lenient().when(config.getReplication()).thenReturn(replicationConfig);

        requestObserver = new RequestObserver();
        transactionManager = spy(new DefaultTransactionManager(config, eventPublisher,
                executionContextRule.getExecutionManager(), logicalClock, mock(MeshClient.class), meshManager,
                repositoryAttributeManager, mock(RepositoryHashService.class), mock(RepositoryManager.class),
                mock(SidebandManager.class), writeTokenRegistry));
        doAnswer(invocation -> {
            ReplicaStateObservedEvent event = invocation.getArgument(0);
            if (event.getRepositoryId().equals(REPOSITORY_ID)) {
                publishedState = event.getObservation().getState();
            }
            return null;
        }).when(eventPublisher).publish(any(ReplicaStateObservedEvent.class));

        leader = new RemoteTransactionLeader(requestObserver, builderFactory, eventPublisher,
                executionContextRule.getExecutionManager(), ForkJoinPool.commonPool(), logicalClock, meshManager,
                replicaStateRegistry, repositoryAttributeManager, repositoryHashService, repositoryManager,
                transactionExecutor, transactionManager, writeTokenRegistry);

        rpcTransaction = RpcTransaction.newBuilder()
                .setId(TRANSACTION_ID)
                .setRepository(REPOSITORY_ID)
                .setCreateRepository(RpcCreateRepositoryRequest.getDefaultInstance())
                .build();
        doAnswer(invocation -> {
            RpcTransaction.Builder builder = invocation.getArgument(0);
            return builder.setCreateRepository(rpcTransaction.getCreateRepository());
        }).when(localOperation).populateTransaction(any());
        when(localOperation.getRepository()).thenReturn(repository);
        //Mock that the write token is available immediately by default. Selected tests will set this to false.
        writeTokenAvailable = true;
        when(writeTokenRegistry.acquireToken(same(repository), anyString(), anyLong(), any()))
                .thenAnswer(invocation -> {
                    writeTokenHandler = invocation.getArgument(3);
                    if (writeTokenAvailable) {
                        writeTokenHandler.onGranted(mock(WriteToken.class));
                    }
                    return writeTokenPromise;
                });
        when(repositoryHashService.getContentHash(repository)).thenReturn(CONTENT_HASH);
        when(repositoryAttributeManager.incrementAndGetVersion(REPOSITORY_ID)).thenReturn(VERSION + 1L);
    }

    @Test(timeout = TIMEOUT_MS)
    public void testFailureToRollBackSendsErrorToLeader() {
        executionContextRule.runInNewContext(() -> {
            //Mock a local operation that prepares to commit and is instructed to roll back
            mockSuccessfulPrepare();
            doThrow(new IllegalStateException("Simulating rollback failure")).when(localOperation).rollBack();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Send the request to roll back, which should unblock the local operation and trigger
            //localOperation.rollBack when the transaction is closed, resulting in a rollback failed
            //being sent to the leader
            sendRollBackRequest();
            //Verify a rollback failure was sent as a normal response, not as an error
            awaitRollBackResponse(false);
            sendSummary(false, LOCAL_ID);
            verifyTokenReleased();

            //Verify that we've ended up in the ROLLED_BACK state
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testFailureToCommitSendsCommitFailed() {
        executionContextRule.runInNewContext(() -> {
            //Mock a local operation that prepares to commit
            mockSuccessfulPrepare();
            doThrow(new IllegalStateException("Simulating commit failure")).when(localOperation).commit();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Send the request to commit, which should unblock the local operation and will trigger
            //localOperation.commit when the transaction is closed, resulting in a commit failed
            //being sent to the leader
            sendCommitRequest();
            //Verify a commit failure was sent as a normal response, not as an error
            awaitCommitResponse(false);
            sendSummary(true, LOCAL_ID);
            verifyTokenReleased();

            //Verify that we've ended up in the ROLLED_BACK state
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testFailureAfterPreparedToCommitSendsCommitFailed() {
        executionContextRule.runInNewContext(() -> {
            //Mock a local operation that prepares, but then throws before the transaction can commit
            mockLocalOperation(() -> {
                PrepareResult result = transactionManager.prepared(localOperation);
                assertTrue("Transaction leader voted NO but expected YES", result.isSuccess());

                throw new IllegalStateException("Simulation exception COMMITTING -> COMMITTED");
            });

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Send the request to commit, which should unblock the local operation and result in the
            //IllegalStateException being thrown, simulating a failure to commit.
            sendCommitRequest();
            //Verify a commit failure was sent as a normal response, not as an error
            awaitCommitResponse(false);
            sendSummary(true, LOCAL_ID); //true because the leader instruction was to commit
            verifyTokenReleased();

            //Verify that we've ended up in the ROLLED_BACK state
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testExceptionAfterPreparedToRollBackSendsRollbackFailed() {
        executionContextRule.runInNewContext(() -> {
            //Mock a local operation that prepares, but then throws before the transaction can roll back
            mockLocalOperation(() -> {
                PrepareResult result = transactionManager.prepared(localOperation);
                assertFalse("Transaction leader voted YES but expected NO", result.isSuccess());

                throw new CreateRepositoryFailedException(REPOSITORY_ID, "Simulating expected operation exception");
            });

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Send the request to roll back, which should unblock the local operation and result in the
            //IllegalStateException being thrown, simulating a rollback failure.
            sendRollBackRequest();
            awaitRollBackResponse(true);
            sendSummary(false);
            verifyTokenReleased();

            //Verify that we've ended up in the ROLLED_BACK state
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyNoEventPublished();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testLeaderErrorTriggersRollback() {
        executionContextRule.runInNewContext(() -> {
            //mock a local operation that prepares to commit
            mockSuccessfulPrepare();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Simulate the remote leader sending an error
            leader.onError(Status.INTERNAL
                    .withDescription("Simulating leader error")
                    .asRuntimeException());

            //This should release the operation, trigger a roll back and allow it to complete
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            //Because the leader just bailed on us, the replica should be observed as UNKNOWN
            verifyReplicaStateObservedEventPublished(ReplicaState.UNKNOWN);
            verifyTokenReleased();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testLocalFailureFollowedByLeaderRollback() {
        executionContextRule.runInNewContext(() -> {
            mockLocalOperation(() -> {
                throw new CreateRepositoryFailedException(REPOSITORY_ID, "Simulated operation failure");
            });

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.NO, localVote);
            //The write token should not be acquired or released
            verifyNoInteractions(writeTokenRegistry);

            //If the local operation fails to prepare, the leader should not send a commit or rollback
            //request. Instead, the summary is used to indicate whether the overall transaction committed
            //anyway (a quorum was reached) or was rolled back
            sendSummary(false);

            //Verify that we've successfully rolled back
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyNoEventPublished();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testLocalTokenNotImmediatelyAvailable() {
        executionContextRule.runInNewContext(() -> {
            //Don't grant the token when the prepared callback is returned
            writeTokenAvailable = false;

            //Mock a successful local prepare
            mockSuccessfulPrepare();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVoteAndVerifyTokenGranted(false);
            assertEquals(Vote.YES, localVote);
            //We shouldn't have recalculated the content hash yet
            verify(repositoryHashService, never()).maybeRecalculateContentHash(repository);
            //Verify that the token has been requested (but not yet granted)
            verify(writeTokenRegistry).acquireToken(same(repository), anyString(), eq(TIME_PREPARE), any());

            //Have the token become available and verify that it is forwarded to the leader
            WriteToken token = mock(WriteToken.class);
            writeTokenHandler.onGranted(token);
            //We've sent over the token, we should have recalculated the content hash
            verify(repositoryHashService).maybeRecalculateContentHash(repository);
            awaitTokenGrant();

            //Now simulate the registry requesting the token to be yielded and verify that the request is forwarded
            writeTokenHandler.yieldToken();
            awaitYieldTokenRequested();
            //But the local token should not be released yet
            verify(token, never()).yieldToken();

            //Have the leader yield the token
            sendTokenYielded();
            verify(token).yieldToken();

            //Now have the leader send a commit instruction, even though it does not hold this node's token
            sendCommitRequest();
            awaitCommitResponse(true);

            sendSummary(true, LOCAL_ID);
            //Verify the token request has been canceled
            verifyTokenReleased();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareAbstainsForInconsistentNonexistentRepositoryWhenNotCreating() {
        //Replace the create request with something else
        rpcTransaction = rpcTransaction.toBuilder()
                .setUpdateRefs(RpcUpdateRefsRequest.getDefaultInstance())
                .build();

        reset(repositoryManager); //Make findById return null again
        when(replicaStateRegistry.getState(eq(REPOSITORY_ID))).thenReturn(ReplicaState.INCONSISTENT);

        executionContextRule.runInNewContext(() -> {
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.ABSTAIN, localVote);

            sendSummary(true, LOCAL_ID);

            //We can't use assertTransactionCompleted here because no transaction is run
            assertTrue("The request was not completed", requestObserver.isCompleted());
            //We can't use verifyReplicaStateObservedEventPublished because it uses assertTransactionCompleted
            assertEquals("No event was published", ReplicaState.INCONSISTENT, publishedState);
        });

        //Verify we never actually start a transaction or try to execute the update
        verify(transactionExecutor, never()).execute(any());
        verify(transactionManager, never()).start(any());
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareAbstainsForInconsistentRepositoryBeingMigratedToLocalNode() {
        //Replace the create request with something else
        rpcTransaction = rpcTransaction.toBuilder()
                .setUpdateRefs(RpcUpdateRefsRequest.getDefaultInstance())
                .build();

        //Mock the repository's partition replica being under migration to the local node
        MeshNode other = mock(MeshNode.class);
        when(meshManager.getPartitionMigrations(PARTITION_ID))
                .thenReturn(ImmutableSet.of(new PartitionReplicaMigration.Builder()
                        .partition(PARTITION_ID)
                        .sourceNode(other)
                        .targetNode(localNode)
                        .build()));
        when(replicaStateRegistry.getState(eq(REPOSITORY_ID))).thenReturn(ReplicaState.REPAIRING);

        executionContextRule.runInNewContext(() -> {
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.ABSTAIN, localVote);

            sendSummary(true, LOCAL_ID);

            //We can't use assertTransactionCompleted here because no transaction is run
            assertTrue("The request was not completed", requestObserver.isCompleted());
            //We can't use verifyReplicaStateObservedEventPublished because it uses assertTransactionCompleted
            assertEquals("No event was published", ReplicaState.INCONSISTENT, publishedState);
        });

        //Verify we never actually start a transaction or try to execute the update
        verify(transactionExecutor, never()).execute(any());
        verify(transactionManager, never()).start(any());
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareRejectsForInconsistentRepositoryBeingMigratedFromLocalNode() {
        //Replace the create request with something else
        rpcTransaction = rpcTransaction.toBuilder()
                .setUpdateRefs(RpcUpdateRefsRequest.getDefaultInstance())
                .build();

        //Mock the repository's partition replica being under migration _from_ the local node. In this case,
        //the follower should _not_ abstain.
        MeshNode other = mock(MeshNode.class);
        when(meshManager.getPartitionMigrations(PARTITION_ID))
                .thenReturn(ImmutableSet.of(new PartitionReplicaMigration.Builder()
                        .partition(PARTITION_ID)
                        .sourceNode(localNode)
                        .targetNode(other)
                        .build()));
        when(replicaStateRegistry.getState(eq(REPOSITORY_ID))).thenReturn(ReplicaState.REPAIRING);

        executionContextRule.runInNewContext(() -> {
            mockLocalOperation(() -> {
                //Simulate the update refs operation failing because the operation fails to prepare
                throw new CommandFailedException("Simulating process failure", new ProcessSummary.Builder("cmd-line")
                        .exitCode(1)
                        .build());
            });

            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.NO, localVote);

            sendSummary(true, LOCAL_ID);

            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareHandlesMultipleRequestFragmentsWhenAbstaining() {
        //This test verifies that RemoteTransactionLeader ignores secondary RpcPrepareTransactionRequest fragments
        //when abstaining as a result of the first request fragment.
        executionContextRule.runInNewContext(() -> {
            //Prepare two ref updates, which we'll send in 2 different fragments
            RpcProposedRefUpdate refUpdate1 = RpcProposedRefUpdate.newBuilder()
                    .setNewHash("1234567890abcdef123451234567890abcdef12345")
                    .setOldHash("abcdef123451234567890abcdef123451234567890")
                    .setRefId(ByteString.copyFromUtf8("refs/heads/branch-1"))
                    .build();
            RpcProposedRefUpdate refUpdate2 = RpcProposedRefUpdate.newBuilder()
                    .setNewHash("0000567890abcdef123451234567890abcdef12345")
                    .setOldHash("abcdef123451234567890abcdef123451234560000")
                    .setRefId(ByteString.copyFromUtf8("refs/heads/branch-2"))
                    .build();

            //Simulate an inconsistent local replica to trigger an ABSTAIN vote
            reset(repositoryManager); //Make findById return null again
            when(replicaStateRegistry.getState(eq(REPOSITORY_ID))).thenReturn(ReplicaState.INCONSISTENT);

            //Send a the first prepare request fragment
            leader.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                            .setLeader(LEADER_ID)
                            .setTimeout(2000L)
                            .setTransaction(RpcTransaction.newBuilder()
                                    .setId(TRANSACTION_ID)
                                    .setRepository(REPOSITORY_ID)
                                    .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                                            .setRepository(REPOSITORY_ID)
                                            .setTransactional(true)
                                            .addRefUpdates(refUpdate1))
                                    .setMoreDetails(true)))
                    .setLogicalTimestamp(TIME_PREPARE)
                    .build());

            //Verify that an ABSTAIN vote is immediately sent, not waiting for the last request fragment
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.ABSTAIN, localVote);

            //Send another prepare request fragment with the second ref update
            leader.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                            .setLeader(LEADER_ID)
                            .setTimeout(2000L)
                            .setTransaction(RpcTransaction.newBuilder()
                                    .setId(TRANSACTION_ID)
                                    .setRepository(REPOSITORY_ID)
                                    .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                                            .setRepository(REPOSITORY_ID)
                                            .setTransactional(true)
                                            .addRefUpdates(refUpdate2))))
                    .setLogicalTimestamp(TIME_PREPARE + 1L)
                    .build());

            //Verify that the second fragment was ignored and did not trigger the transaction to be replicated after all
            verifyNoInteractions(transactionExecutor);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareParticipatesWithNonexistentRepositoryIfLocallyConsistent() {
        //Replace the create request with something else
        rpcTransaction = rpcTransaction.toBuilder()
                .setUpdateRefs(RpcUpdateRefsRequest.getDefaultInstance())
                .build();

        reset(repositoryManager); //Make findById return null again
        when(replicaStateRegistry.getState(eq(REPOSITORY_ID))).thenReturn(ReplicaState.CONSISTENT);

        executionContextRule.runInNewContext(() -> {
            mockLocalOperation(() -> {
                //Simulate the update refs operation failing because the repository doesn't exist
                throw new RepositoryNotFoundException(REPOSITORY_ID, null);
            });

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.NO, localVote);
            //The write token should not be acquired or released
            verifyNoInteractions(writeTokenRegistry);

            //Depending on the qourum mode, the leader could still decide to commit. In that case, it
            //shouldn't send us a commit (or rollback) instruction; it should simply send us a summary
            //indicating the transaction was committed and we're now inconsistent
            sendSummary(true, LOCAL_ID);

            //Verify that we've successfully rolled back
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPrepareParticipatesWithNonexistentRepositoryWhenCreating() {
        reset(repositoryManager); //Make findById return null again

        executionContextRule.runInNewContext(() -> {
            mockSuccessfulPrepare();

            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            sendCommitRequest();
            awaitCommitResponse(true);
            sendSummary(true);
            verifyTokenReleased();

            assertTransactionCompleted(TransactionState.COMMITTED);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPreparedLocallyAndCommitFails() {
        executionContextRule.runInNewContext(() -> {
            mockSuccessfulPrepare();
            doThrow(RuntimeException.class).when(localOperation).commit();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Instruct the local operation to commit, which should trigger our trained exception and fail
            sendCommitRequest();
            awaitCommitResponse(false);
            sendSummary(true, LOCAL_ID);
            verifyTokenReleased();

            //Verify that we're rolled back
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            //Verify that the replica has been marked as inconsistent
            verifyReplicaStateObservedEventPublished(ReplicaState.INCONSISTENT);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPreparedLocallyAndRemoteLeaderVotesToRollBack() {
        executionContextRule.runInNewContext(() -> {
            mockLocalOperation(() -> {
                PrepareResult result = transactionManager.prepared(localOperation);
                assertFalse("Transaction leader voted YES but expected NO", result.isSuccess());
            });

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Instruct the local operation to roll back and wait for it to succeed
            sendRollBackRequest();
            awaitRollBackResponse(true);
            sendSummary(false);
            verifyTokenReleased();

            //Verify that we've successfully rolled back
            assertTransactionCompleted(TransactionState.ROLLED_BACK);
            verifyNoEventPublished();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPreparedLocallyAndRemoteLeaderVotesToCommit() {
        executionContextRule.runInNewContext(() -> {
            mockSuccessfulPrepare();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            //We should not have recalculated the content hash yet
            verify(repositoryHashService, never()).maybeRecalculateContentHash(repository);
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();
            //We've sent over the token, we should have recalculated the content hash
            verify(repositoryHashService).maybeRecalculateContentHash(repository);

            //Instruct the local operation to commit and wait for it to succeed
            sendCommitRequest();
            awaitCommitResponse(true);
            sendSummary(true);
            verifyTokenReleased();

            //Verify that we've successfully committed
            assertTransactionCompleted(TransactionState.COMMITTED);
            verifyNoEventPublished();
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testPreparedLocallyAndRemoteLeaderVotesToCommitBeforePrepared() {
        executionContextRule.runInNewContext(() -> {
            CountDownLatch latch = new CountDownLatch(2);

            mockLocalOperation(() -> {
                latch.countDown();
                try {
                    latch.await();
                } catch (InterruptedException interruptedException) {
                    Thread.currentThread().interrupt();
                    throw new RuntimeException("interrupted");
                }
                PrepareResult result = transactionManager.prepared(localOperation);
                assertTrue("Transaction leader voted NO but expected YES", result.isSuccess());
            });

            //Send the prepare request, then the commit request and wait for the local operation to be ready to prepare
            sendPrepareRequest();
            //Token should not be acquired yet
            verifyNoInteractions(writeTokenRegistry);
            sendCommitRequest();

            latch.countDown();
            try {
                latch.await();
            } catch (InterruptedException interruptedException) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("interrupted");
            }

            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);
            verifyTokenAcquired();

            //Instruct the local operation to commit and wait for it to succeed
            awaitCommitResponse(true);
            sendSummary(true);
            verifyTokenReleased();

            //Verify that we've successfully committed
            assertTransactionCompleted(TransactionState.COMMITTED);
            verifyNoEventPublished();
        });
    }
    
    @Test(timeout = TIMEOUT_MS)
    public void testPrepareWaitsForLastPrepareRequestFragment() {
        //This test verifies that RemoteTransactionLeader merges RpcPrepareTransactionRequest fragments
        //for ref-updates that are so large that they are split across multiple messages.
        executionContextRule.runInNewContext(() -> {
            //Prepare two ref updates, which we'll send in 2 different fragments
            RpcProposedRefUpdate refUpdate1 = RpcProposedRefUpdate.newBuilder()
                    .setNewHash("1234567890abcdef123451234567890abcdef12345")
                    .setOldHash("abcdef123451234567890abcdef123451234567890")
                    .setRefId(ByteString.copyFromUtf8("refs/heads/branch-1"))
                    .build();
            RpcProposedRefUpdate refUpdate2 = RpcProposedRefUpdate.newBuilder()
                    .setNewHash("0000567890abcdef123451234567890abcdef12345")
                    .setOldHash("abcdef123451234567890abcdef123451234560000")
                    .setRefId(ByteString.copyFromUtf8("refs/heads/branch-2"))
                    .build();

            //We expect the RemoteTransactionLeader to merge the RpcTransaction fragments into a single
            //transaction with all the ref updates.
            RpcTransaction expectedTransaction = RpcTransaction.newBuilder()
                    .setId(TRANSACTION_ID)
                    .setRepository(REPOSITORY_ID)
                    .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                            .setRepository(REPOSITORY_ID)
                            .setTransactional(true)
                            .addRefUpdates(refUpdate1)
                            .addRefUpdates(refUpdate2))
                    .build();

            //Simulate that local transaction successfully replaying the operations
            doAnswer(invocation -> {
                RpcTransaction.Builder builder = invocation.getArgument(0);
                return builder.setUpdateRefs(expectedTransaction.getUpdateRefs());
            }).when(localOperation).populateTransaction(any());
            mockLocalOperation(expectedTransaction, () -> {
                PrepareResult result = transactionManager.prepared(localOperation);
                assertTrue("Transaction leader voted NO but expected YES", result.isSuccess());
            });

            //Send a the first prepare request fragment
            leader.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                            .setLeader(LEADER_ID)
                            .setTimeout(2000L)
                            .setTransaction(RpcTransaction.newBuilder()
                                    .setId(TRANSACTION_ID)
                                    .setRepository(REPOSITORY_ID)
                                    .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                                            .setRepository(REPOSITORY_ID)
                                            .setTransactional(true)
                                            .addRefUpdates(refUpdate1))
                                    .setMoreDetails(true)))
                    .setLogicalTimestamp(TIME_PREPARE)
                    .build());

            //Send another prepare request fragment with the second ref update
            leader.onNext(RpcTransactLeaderFragment.newBuilder()
                    .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                            .setLeader(LEADER_ID)
                            .setTimeout(2000L)
                            .setTransaction(RpcTransaction.newBuilder()
                                    .setId(TRANSACTION_ID)
                                    .setRepository(REPOSITORY_ID)
                                    .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                                            .setRepository(REPOSITORY_ID)
                                            .setTransactional(true)
                                            .addRefUpdates(refUpdate2))))
                    .setLogicalTimestamp(TIME_PREPARE + 1L)
                    .build());

            //Verify that the local operation is prepared
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);

            //Verify that the combined RpcUpdateRefsRequest was passed to the RpcTransactionExecutor
            verify(transactionExecutor).execute(expectedTransaction);
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testStagedPackIsUsedIfProvided() {
        executionContextRule.runInNewContext(() -> {
            mockLocalOperation(() -> {
                PrepareResult result = transactionManager.prepared(localOperation);
                assertTrue("Transaction leader voted NO but expected YES", result.isSuccess());
            });
            Future<Path> indexPackFuture = mockIndexPack();

            //Send a staged pack before calling prepare (a simulated push)
            sendPack();

            //Send the prepare request and wait for the local operation to prepare
            sendPrepareRequest();
            Vote localVote = awaitLocalVote();
            assertEquals(Vote.YES, localVote);

            //Verify that the follower transaction waits for index-pack to complete before preparing its transaction,
            //and that the staging directory is passed as an alternate to transactionExecutor
            InOrder inOrder = inOrder(indexPackFuture, transactionExecutor);
            try {
                inOrder.verify(indexPackFuture).get(anyLong(), any());
            } catch (InterruptedException | ExecutionException | TimeoutException e) {
                fail("verify failed");
            }
            inOrder.verify(transactionExecutor).execute(any());

            //The pack is written to a staging directory, under <repo-dir>/app-info/tx/transaction-id
            Path incomingDir = leader.getIncomingDir().orElse(null);
            Path expected = RemoteTransactionLeader.getIncomingDir(repository, TRANSACTION_ID);
            assertEquals(expected, incomingDir);

            //Write something to the staging area and validate that it is published on commit
            try {
                Path packDir = MoreFiles.mkdir(incomingDir.resolve("pack"));
                MoreFiles.write(packDir.resolve("pack-abcdef.pack"), "testing pack");
                MoreFiles.write(packDir.resolve("pack-abcdef.idx"), "testing index");

                //Register a transaction listener to check whether the pack is published before commit.
                MutableBoolean published = new MutableBoolean();
                TransactionListener txListener = new TransactionListener() {

                    @Override
                    public void beforeCommit(@Nonnull Transaction transaction) {
                        try {
                            //Verify the pack has been published
                            Path repoPackDir = repository.getPath().resolve("objects/pack");
                            assertEquals("testing pack", MoreFiles.toString(repoPackDir.resolve("pack-abcdef.pack")));
                            assertEquals("testing index", MoreFiles.toString(repoPackDir.resolve("pack-abcdef.idx")));

                            //Verify the staging area is cleaned up after the staged pack is published
                            assertFalse(MoreFiles.dirExists(incomingDir));
                            //Verify the leader no longer returns an incomingDir, because the directory no longer exists
                            assertNull(leader.getIncomingDir().orElse(null));

                            published.setTrue();
                        } catch (Throwable t) {
                            //Exceptions thrown from beforeCommit don't propagate, so for us to be able to see
                            //the error that happened we need to write it somewhere
                            System.err.println("beforeCommit validation failed");
                            t.printStackTrace(System.err);
                        }
                    }
                };

                Transaction transaction = transactionManager.getCurrentTransaction().orElse(null);
                assertNotNull("A transaction should be active", transaction);
                transaction.addListener(txListener);

                sendCommitRequest();
                awaitCommitResponse(true);

                //Verify that the transaction was actually committed and that our validation did run
                assertTrue("beforeCommit validation either did not run or failed", published.isTrue());

                sendSummary(true);
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testStagedPackIsPublishedIfLocalTransactionFailsButLeaderCommits() {
        executionContextRule.runInNewContext(() -> {
            //Mock the local prepare operation throwing an exception without triggering tx callbacks
            //(e.g. a .lock file preventing a ref update).
            mockLocalOperation(() -> {
                throw new CommandFailedException("mocking failure", new ProcessSummary.Builder("command-line")
                        .exitCode(1)
                        .build());
            });
            mockIndexPack();

            //Send a staged pack before calling prepare (a simulated push)
            sendPack();

            //Send the prepare request and wait for the local operation to prepare (and fail)
            sendPrepareRequest();

            Vote localVote = awaitLocalVote();
            assertEquals(Vote.NO, localVote);

            //The pack is written to a staging directory, under <repo-dir>/app-info/tx/transaction-id
            Path incomingDir = leader.getIncomingDir().orElse(null);
            Path expected = RemoteTransactionLeader.getIncomingDir(repository, TRANSACTION_ID);
            assertEquals(expected, incomingDir);

            //Write something to the staging area and validate that it is published when the summary is sent
            try {
                Path packDir = MoreFiles.mkdir(incomingDir.resolve("pack"));
                MoreFiles.write(packDir.resolve("pack-abcdef.pack"), "testing pack");
                MoreFiles.write(packDir.resolve("pack-abcdef.idx"), "testing index");
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }

            //Have the leader send a summary indicating that the transaction committed and that the local node
            //is now inconsistent.
            sendSummary(true, LOCAL_ID);

            //Verify that the staged pack has been published
            assertTrue(MoreFiles.canOpen(MoreFiles.resolve(repository.getPath(), "objects", "pack", "pack-abcdef.pack")));
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testStagedPackIsCleanedUpOnError() {
        executionContextRule.runInNewContext(() -> {
            mockIndexPack();

            //Send a staged pack before calling prepare (a simulated push)
            sendPack();

            //Verify that index-pack has been started and that the staging directory has been created
            verify(builderFactory.builder(repository)).indexPack();
            Path incomingDir = RemoteTransactionLeader.getIncomingDir(repository, TRANSACTION_ID);
            assertTrue(MoreFiles.dirExists(incomingDir));

            //Simulate an error
            leader.onError(Status.UNKNOWN.asRuntimeException());

            //Verify that the stagingDir is cleaned up
            assertFalse(MoreFiles.dirExists(incomingDir));
        });
    }

    @Test(timeout = TIMEOUT_MS)
    public void testStagedPackIsRejectedIfRepositoryDoesNotExist() {
        executionContextRule.runInNewContext(() -> {
            when(repositoryManager.findById(REPOSITORY_ID)).thenReturn(null); //not found

            //Send a staged pack before calling prepare (a simulated push)
            sendPack();

            awaitPackAbort();
        });
    }

    @Test
    public void testSummaryWithoutPrepare() {
        executionContextRule.runInNewContext(() -> {
            //Send a staged pack
            mockIndexPack();
            sendPack();

            //Verify that index-pack has been started and that the staging directory has been created
            verify(builderFactory.builder(repository)).indexPack();
            Path incomingDir = RemoteTransactionLeader.getIncomingDir(repository, TRANSACTION_ID);
            assertTrue(MoreFiles.dirExists(incomingDir));

            //Simulate the leader wrapping up the transaction with a summary without ever sending any
            //prepare, commit or rollback fragments. The local ID is marked inconsistent here to test
            //whether inconsistent nodes are ignored (as they should be)
            sendSummary(false, LOCAL_ID);

            //Verify that the stagingDir is cleaned up
            assertFalse(MoreFiles.dirExists(incomingDir));
            //Verify the reported inconsistency for the local node is not published
            verify(eventPublisher, never()).publish(isA(ReplicaStateObservedEvent.class));
        });
    }

    private void assertTransactionCompleted(TransactionState state) {
        assertEquals(state, awaitTransactionCompleted());
        assertTrue("The request was not completed", requestObserver.isCompleted());
    }

    private void awaitCommitResponse(boolean committed) {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasCommitResponse());
        assertNotEquals("Should have sent a commit" + (committed ? "ted" : " failed") + " response",
                committed, fragment.getCommitResponse().hasCommitFailed());

        if (committed) {
            assertTrue("Should have included a replicaVersion", fragment.getCommitResponse().hasReplicaVersion());
            RpcReplicaVersion replicaVersion = fragment.getCommitResponse().getReplicaVersion();
            assertEquals(CONTENT_HASH, replicaVersion.getContentHash());
            assertEquals(VERSION + 1, replicaVersion.getVersion());
        }
    }

    private Vote awaitLocalVote() {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasPrepareResponse());

        return fragment.getPrepareResponse().getVote();
    }

    private Vote awaitLocalVoteAndVerifyTokenGranted(boolean granted) {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasPrepareResponse());

        RpcPrepareTransactionResponse prepareResponse = fragment.getPrepareResponse();
        assertEquals(granted, prepareResponse.hasWriteTokenGranted());
        if (granted) {
            RpcReplicaVersion replicaVersion = prepareResponse.getWriteTokenGranted().getReplicaVersion();
            assertEquals(CONTENT_HASH, replicaVersion.getContentHash());
            assertEquals(VERSION, replicaVersion.getVersion());
        }

        return prepareResponse.getVote();
    }

    private void awaitPackAbort() {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue("Should have sent a request to abort the pack", fragment.hasAbortStagingPack());
    }

    private void awaitRollBackResponse(boolean rolledBack) {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasRollBackResponse());
        assertNotEquals("Should have sent a roll" + (rolledBack ? "ed back" : "back failed") + " response",
                rolledBack, fragment.getRollBackResponse().hasRollbackFailed());
    }

    private void awaitTokenGrant() {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasTokenGranted());
    }

    private TransactionState awaitTransactionCompleted() {
        try {
            return completionState.get(1, TimeUnit.SECONDS);
        } catch (InterruptedException interruptedException) {
            throw new RuntimeException("Interrupted while waiting for transaction to complete");
        } catch (ExecutionException executionException) {
            Throwable cause = executionException.getCause();
            Throwables.throwIfUnchecked(cause);
            throw new RuntimeException("Transaction failed", cause);
        } catch (TimeoutException e) {
            throw new RuntimeException("Timed out waiting for transaction to complete");
        }
    }

    private void awaitYieldTokenRequested() {
        RpcTransactFollowerFragment fragment = requestObserver.next();
        assertNotNull(fragment);
        assertTrue(fragment.hasYieldTokenRequest());
    }

    private Future<Path> mockIndexPack() {
        Future<Path> future = mock(Future.class);
        GitIndexPackBuilder builder = mock(GitIndexPackBuilder.class, Answers.RETURNS_SELF);
        GitProcess<Path> process = mock(GitProcess.class);
        ArgumentCaptor<IndexPackStdioHandler> handlerCaptor = ArgumentCaptor.forClass(IndexPackStdioHandler.class);
        when(builder.build(handlerCaptor.capture())).thenReturn(process);
        when(builderFactory.builder(any(Repository.class)).indexPack()).thenReturn(builder);
        when(process.start()).thenAnswer(invocation -> {
            IndexPackStdioHandler handler = handlerCaptor.getValue();
            if (handler != null) {
                OsProcess osProcess = mock(OsProcess.class);
                handler.onPreStart(osProcess);
                handler.onStart(osProcess);
            }
            return future;
        });
        return future;
    }

    private void mockLocalOperation(Runnable operation) {
        mockLocalOperation(rpcTransaction, operation);
    }

    private void mockLocalOperation(RpcTransaction rpcTransaction, Runnable operation) {
        doAnswer(invocation -> {
            //Make the current transaction visible to the test methods
            Transaction transaction = transactionManager.getCurrentTransaction().orElse(null);
            assertNotNull(transaction);
            transaction.addListener(new TransactionListener() {

                @Override
                public void afterCompletion(@Nonnull Transaction transaction) {
                    completionState.complete(transaction.getState());
                }
            });

            operation.run();
            return null;
        }).when(transactionExecutor).execute(eq(rpcTransaction));
    }

    private void mockSuccessfulPrepare() {
        mockLocalOperation(() -> {
            PrepareResult result = transactionManager.prepared(localOperation);
            assertTrue("Transaction leader voted NO but expected YES", result.isSuccess());
        });
    }

    private void sendRollBackRequest() {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setRollBackRequest(RpcRollBackTransactionRequest.getDefaultInstance())
                .setLogicalTimestamp(TIME_ROLLBACK)
                .build());
    }

    private void sendCommitRequest() {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setCommitRequest(RpcCommitTransactionRequest.newBuilder()
                        .setVersion(VERSION))
                .setLogicalTimestamp(TIME_COMMIT)
                .build());
    }

    private void sendPack() {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setStagePackRequest(RpcStagePackRequest.newBuilder()
                        .setRepository(REPOSITORY_ID)
                        .setTransactionId(TRANSACTION_ID))
                .setLogicalTimestamp(TIME_PACK)
                .build());
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setPackChunk(RpcStreamChunk.newBuilder()
                        .setData(ByteString.copyFromUtf8("chunk1")))
                .setLogicalTimestamp(TIME_PACK)
                .build());
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setPackChunk(RpcStreamChunk.newBuilder()
                        .setData(ByteString.copyFromUtf8("chunk2"))
                        .setClosed(true))
                .setLogicalTimestamp(TIME_PACK)
                .build());
    }

    private void sendPrepareRequest() {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setPrepareRequest(RpcPrepareTransactionRequest.newBuilder()
                        .setLeader(LEADER_ID)
                        .setTimeout(2000L)
                        .setTransaction(rpcTransaction))
                .setLogicalTimestamp(TIME_PREPARE)
                .build());
    }

    private void sendSummary(boolean committed, String... inconsistentReplicas) {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setSummary(RpcTransactionSummary.newBuilder()
                        .addAllInconsistentReplicas(Arrays.asList(inconsistentReplicas))
                        .setCommitted(committed))
                .setLogicalTimestamp(TIME_SUMMARY)
                .build());
    }

    private void sendTokenYielded() {
        leader.onNext(RpcTransactLeaderFragment.newBuilder()
                .setTokenYielded(RpcWriteTokenYielded.getDefaultInstance())
                .build());
    }

    private void verifyReplicaStateObservedEventPublished(ReplicaState expectedState) {
        //Events are published before the transaction completes
        awaitTransactionCompleted();
        assertEquals("No event was published", expectedState, publishedState);
    }

    private void verifyNoEventPublished() {
        awaitTransactionCompleted();
        assertNull("An event was published but not expected", publishedState);
    }

    private void verifyTokenAcquired() {
        verify(writeTokenRegistry).acquireToken(eq(repository), anyString(), eq(TIME_PREPARE), any());
    }

    private void verifyTokenReleased() {
        verify(writeTokenPromise).close();
    }

    private static class RequestObserver implements StreamObserver<RpcTransactFollowerFragment> {

        private final BlockingQueue<RpcTransactFollowerFragment> fragments = new LinkedBlockingDeque<>(2);

        private boolean completed;
        private Throwable exception;

        @Override
        public void onCompleted() {
            completed = true;
        }

        @Override
        public void onError(Throwable t) {
            exception = t;
            //Add a marker fragment to unblock next()
            enqueue(RpcTransactFollowerFragment.getDefaultInstance());
        }

        @Override
        public void onNext(RpcTransactFollowerFragment value) {
            enqueue(value);
        }

        boolean isCompleted() {
            return completed || exception != null;
        }

        RpcTransactFollowerFragment next() {
            throwIfErrorSentToLeader();
            try {
                RpcTransactFollowerFragment fragment = fragments.take();
                throwIfErrorSentToLeader();
                return fragment;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException(e);
            }
        }

        private void enqueue(RpcTransactFollowerFragment fragment) {
            try {
                fragments.put(fragment);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException(e);
            }
        }

        private void throwIfErrorSentToLeader() {
            if (exception != null) {
                Throwables.throwIfUnchecked(exception);
                throw new RuntimeException(exception);
            }
        }
    }
}