package com.atlassian.bitbucket.mesh.transaction;

import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.client.ReplicateParameters;
import com.atlassian.bitbucket.mesh.config.ApplicationConfig;
import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.config.ReplicationConfig;
import com.atlassian.bitbucket.mesh.coordination.LamportClock;
import com.atlassian.bitbucket.mesh.coordination.LogicalClock;
import com.atlassian.bitbucket.mesh.coordination.WriteLock;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.management.SidebandManager;
import com.atlassian.bitbucket.mesh.network.*;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.protobuf.ByteString;
import org.apache.commons.lang3.mutable.MutableBoolean;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.commons.lang3.tuple.Pair;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.InOrder;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.transaction.PrepareResult.*;
import static com.atlassian.bitbucket.mesh.transaction.QuorumMode.*;
import static java.util.Optional.of;
import static org.apache.commons.lang3.tuple.Pair.of;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class LocalTransactionLeaderTest {

    private static final String CONTENT_HASH = "content-hash";
    private static final MeshNode LOCAL = new MeshNode.Builder("local")
            .local(true)
            .name("local")
            .rpcUrl("local")
            .build();
    private static final String PARTITION = "partition";
    private static final String REPOSITORY_ID = "repository-id";
    private static final long TOKEN_TIMESTAMP = 12319L;
    private static final RpcTransaction TX = RpcTransaction.newBuilder()
            .setId("tx")
            .build();
    private static final Iterable<RpcTransaction> TX_FRAGMENTS = ImmutableList.of(TX);
    private static final long VERSION = 77L;
    private static final ReplicaVersion PREPARED_REPLICA_VERSION = new ReplicaVersion(CONTENT_HASH, VERSION);

    @Rule
    public final LogListener logListener = new LogListener(LocalTransactionLeader.class);

    @Mock
    private ApplicationConfig config;
    @Mock
    private EventPublisher eventPublisher;
    private ExecutorService executorService;
    @Mock
    private GrpcConfig grpcConfig;
    private LogicalClock logicalClock;
    @Mock
    private MeshClient meshClient;
    @Mock
    private MeshManager meshManager;
    @Mock
    private ReplicationConfig replicationConfig;
    @Mock
    private ManagedRepository repository;
    @Mock
    private RepositoryAttributeManager repositoryAttributeManager;
    @Mock
    private RepositoryHashService repositoryHashService;
    @Mock
    private SidebandManager sidebandManager;
    @Mock
    private WriteLock writeLock;
    private CompletableFuture<Void> writeLockFuture;
    @Mock
    private WriteTokenRegistry writeTokenRegistry;

    @Before
    public void setUp() {
        logicalClock = new LamportClock();
        //Have the write lock complete immediately by default
        writeLockFuture = CompletableFuture.completedFuture(null);
        when(config.getGrpc()).thenReturn(grpcConfig);
        when(config.getReplication()).thenReturn(replicationConfig);
        when(grpcConfig.getMaxFragmentSize()).thenReturn(64 * 1024L);
        when(meshManager.getLocalNode()).thenReturn(LOCAL);
        when(replicationConfig.getFinalizeInterval()).thenReturn(Duration.ofMinutes(5L));
        when(repository.getId()).thenReturn(REPOSITORY_ID);
        when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(VERSION);
        when(repositoryAttributeManager.incrementAndGetVersion(REPOSITORY_ID)).thenReturn(VERSION + 1);
        when(repositoryHashService.getContentHash(any())).thenReturn(CONTENT_HASH);
        when(writeLock.getLogicalTimestamp()).thenReturn(TOKEN_TIMESTAMP);
        when(writeLock.asCompletionStage()).thenAnswer(invocation -> writeLockFuture);
        //honor any token yield request as long as the write lock hasn't been tripped yet
        when(writeLock.yieldToken(any())).thenAnswer(invocation -> !writeLockFuture.isDone());
        when(writeTokenRegistry.startLock(any(), anyString(), anyCollection())).thenReturn(writeLock);

        executorService = Executors.newFixedThreadPool(1);
    }

    @After
    public void tearDown() throws Exception {
        executorService.shutdownNow();
        executorService.awaitTermination(1, TimeUnit.SECONDS);
    }

    @Test
    public void testGetNode() {
        assertSame(LOCAL, createLeader(UNANIMOUS).getNode());
    }

    @Test
    public void testOnCommittedAlwaysMode() {
        doTestOnCommitted(ALWAYS);
    }

    @Test
    public void testOnCommittedHandlesRemoteFailuresAlwaysMode() {
        doTestOnCommittedHandlesRemoteFailures(ALWAYS);
    }

    @Test
    public void testOnCommittedHandlesRemoteFailuresMajorityMode() {
        doTestOnCommittedHandlesRemoteFailures(MAJORITY);
    }

    @Test
    public void testOnCommittedHandlesRemoteFailuresUnanimousMode() {
        doTestOnCommittedHandlesRemoteFailures(UNANIMOUS);
    }

    @Test
    public void testOnCommittedMajorityMode() {
        doTestOnCommitted(MAJORITY);
    }

    @Test
    public void testOnCommittedUnanimousMode() {
        doTestOnCommitted(UNANIMOUS);
    }

    @Test
    public void testOnCommittedWhenAllNodesFail() {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);
        when(replicaTransaction.await(notNull())).thenThrow(TransactionCommitFailedException.class);
        when(replicaTransaction.getReplicaVersion()).thenReturn(Optional.empty());

        LocalTransactionLeader leader = createLeader(UNANIMOUS);
        leader.onPrepared(repository, TX);
        leader.onCommitted(new RuntimeException());

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).commit(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertThat(summary.getInconsistentReplicasList(), empty());
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentAlwaysMode() {
        doTestOnCommittedWhenRemotelyInconsistent(ALWAYS);
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentMajorityMode() {
        doTestOnCommittedWhenRemotelyInconsistent(MAJORITY);
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentUnanimousMode() {
        doTestOnCommittedWhenRemotelyInconsistent(UNANIMOUS);
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentVersionAlwaysMode() {
        doTestOnCommittedWhenRemotelyInconsistentVersion(ALWAYS);
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentVersionMajorityMode() {
        doTestOnCommittedWhenRemotelyInconsistentVersion(MAJORITY);
    }

    @Test
    public void testOnCommittedWhenRemotelyInconsistentVersionUnanimousMode() {
        doTestOnCommittedWhenRemotelyInconsistentVersion(UNANIMOUS);
    }

    @Test
    public void testOnCommittedWhenReplicaAbstainsAlwaysMode() {
        doTestOnCommittedWhenReplicaAbstains(ALWAYS);
    }

    @Test
    public void testOnCommittedWhenReplicaAbstainsMajorityMode() {
        doTestOnCommittedWhenReplicaAbstains(MAJORITY);
    }

    @Test
    public void testOnCommittedWhenReplicaAbstainsUnanimousMode() {
        doTestOnCommittedWhenReplicaAbstains(UNANIMOUS);
    }

    @Test
    public void testOnCommittedWithLocalFailureAlwaysMode() {
        doTestOnCommittedWithLocalFailure(ALWAYS);
    }

    @Test
    public void testOnCommittedWithLocalFailureMajorityMode() {
        doTestOnCommittedWithLocalFailure(MAJORITY);
    }

    @Test
    public void testOnCommittedWithLocalFailureUnanimousMode() {
        doTestOnCommittedWithLocalFailure(UNANIMOUS);
    }

    @Test
    public void testOnPreparedFollowedByYieldTokenAndFailure() throws ExecutionException, InterruptedException {
        //Don't complete the writeLock until we've requested the token back
        writeLockFuture = new CompletableFuture<>();
        ReplicaTransaction replicaTransaction = setUpReplicaTransactions(SUCCESS).get(0);

        LocalTransactionLeader leader = createLeader(MAJORITY);

        //Let the leader prepare; the replicaTransaction will successfully repair and grant its token.
        //The returned future should not be complete yet because we've mocked the writeLock to not complete yet.
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertFalse(future.isDone());

        ArgumentCaptor<ReplicateParameters> paramsCaptor = ArgumentCaptor.forClass(ReplicateParameters.class);
        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock).asCompletionStage();
        verify(meshClient).replicate(any(), paramsCaptor.capture());
        ReplicaTransactionCallback callback = paramsCaptor.getValue().getCallback();

        //Mock the replica requesting (and getting) its token back
        MeshNode replica = replicaTransaction.getNode();
        assertTrue(callback.yieldToken(replica));
        //Next mock the replica transaction failing (e.g. node going down)
        callback.onReplicaError(replica);

        //Release the write lock and verify that the tx is rolled
        writeLockFuture.complete(null);

        //Verify that the write is rejected (no quorum)
        assertFalse(future.get().isSuccess());

        leader.onRolledBack(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock).asCompletionStage();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(REPOSITORY_ID);
        verify(replicaTransaction, never()).commit(anyLong());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertMarkedInconsistent(summary); //all nodes consistent
    }

    @Test
    public void testOnPreparedHandlesFollowerResponsesWithoutWriteTokensCorrectly() throws ExecutionException, InterruptedException {
        //Two transactions sending SUCCESS without write tokens
        List<MockReplicaTransaction> transactions =
                setUpReplicaTransactions(ImmutableList.of(of(SUCCESS, false), of(SUCCESS, false)));
        mockManagedRepository();
        writeLockFuture = new CompletableFuture<>();

        LocalTransactionLeader leader = createLeader(MAJORITY);
        Future<PrepareResult> future = leader.onPrepared(repository, TX);

        assertFalse(future.isDone());
        verify(writeLock).asCompletionStage();

        //Grant the write tokens after the prepared callback
        transactions.forEach(MockReplicaTransaction::grantWriteToken);

        writeLockFuture.complete(null);

        assertTrue(future.isDone());
        assertTrue(future.get().isSuccess());
    }

    @Test
    public void testOnPreparedSplitsUpUpdateRefsTransaction() {
        int maxSize = 512; // 0.5kb fragments;
        when(grpcConfig.getMaxFragmentSize()).thenReturn((long) maxSize);
        RpcTransaction transaction = RpcTransaction.newBuilder()
                .setId("tx")
                .setRepository("p/0000/h/1234abcdef1234abcdef/r/8729")
                .setUpdateRefs(RpcUpdateRefsRequest.newBuilder()
                        .setRepository("p/0000/h/1234abcdef1234abcdef/r/8729")
                        .setTransactional(true)
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/1")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/2")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/3")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/4")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/5")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/6")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/7")))
                        .addRefUpdates(RpcProposedRefUpdate.newBuilder()
                                .setNewHash("1234512345123451234509876098760987609876")
                                .setOldHash("0987609876098760987612345123451234512345")
                                .setRefId(ByteString.copyFromUtf8("refs/heads/a/long/branch/name/8"))))
                .build();

        assertThat(transaction.getSerializedSize(), greaterThan(maxSize));

        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);
        createLeader(MAJORITY).onPrepared(repository, transaction);

        //Verify that the large fragment has been split into 3 smaller fragments and that each of the fragments
        //is under the maxSize
        ArgumentCaptor<Iterable<RpcTransaction>> fragmentsCaptor = ArgumentCaptor.forClass(Iterable.class);
        verify(replicaTransaction).prepare(fragmentsCaptor.capture(), anyLong());

        List<RpcTransaction> fragments = ImmutableList.copyOf(fragmentsCaptor.getValue());
        assertEquals(3, fragments.size());
        assertThat(fragments.get(0).getSerializedSize(), lessThanOrEqualTo(maxSize));
        assertTrue(fragments.get(0).getMoreDetails());

        assertThat(fragments.get(1).getSerializedSize(), lessThanOrEqualTo(maxSize));
        assertTrue(fragments.get(1).getMoreDetails());

        assertThat(fragments.get(2).getSerializedSize(), lessThanOrEqualTo(maxSize));
        assertFalse(fragments.get(2).getMoreDetails());

        //Verify that the total list of ref updates is identical (and in the same cumulative order)
        List<RpcProposedRefUpdate> cumulativeUpdates = fragments.stream()
                .flatMap(fragment -> fragment.getUpdateRefs().getRefUpdatesList().stream())
                .collect(Collectors.toList());
        assertEquals(transaction.getUpdateRefs().getRefUpdatesList(), cumulativeUpdates);
    }

    @Test
    public void testOnPreparedTakesPartitionMigrationIntoAccountSingleQuorumFailsMajority() throws Exception {
        //Tests a migration scenario where an individual quorum set triggers failure, but the write would be
        //allowed if no special handling for partition migrations was in place.

        //Mock a total of 5 replicas (local, node1, node2, node3, node4) and 2 migrations (node1->node3, node2->node4).
        //In this scenario we'll let node3 and node4 vote NO, which should cause the write to be rejected. Without
        //migration-aware quorum handling, this scenario would result in a commit because 3 out of 5 replicas vote
        //YES which forms a majority
        List<ReplicaTransaction> transactions = setUpReplicaTransactions(SUCCESS, SUCCESS, FAILURE, FAILURE);
        mockManagedRepository();

        PartitionReplicaMigration migration1_3 = new PartitionReplicaMigration.Builder()
                .partition(PARTITION)
                .sourceNode(transactions.get(0).getNode()) //node1
                .targetNode(transactions.get(2).getNode()) //node3
                .build();
        PartitionReplicaMigration migration2_4 = new PartitionReplicaMigration.Builder()
                .partition(PARTITION)
                .sourceNode(transactions.get(1).getNode()) //node2
                .targetNode(transactions.get(3).getNode()) //node4
                .build();
        when(meshManager.getPartitionMigrations("partition")).thenReturn(ImmutableList.of(migration1_3, migration2_4));

        LocalTransactionLeader leader = createLeader(MAJORITY);
        Future<PrepareResult> future = leader.onPrepared(repository, TX);

        //verify that each replica transaction is instructed to rollback
        assertFalse(future.get().isSuccess());
        for (ReplicaTransaction transaction : transactions) {
            verify(transaction).rollBack(anyLong());
        }
    }

    @Test
    public void testOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPassAlways() throws Exception {
        doTestOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPass(ALWAYS);
    }

    @Test
    public void testOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPassMajority() throws Exception {
        doTestOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPass(MAJORITY);
    }

    @Test
    public void testOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPassUnanimous() throws Exception {
        doTestOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPass(UNANIMOUS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionAlways() {
        doTestOnPreparedThenOnCommittedWithNoPartition(ALWAYS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionIdAlways() throws Exception {
        doTestOnPreparedThenOnCommittedWithNoPartitionId(ALWAYS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionIdMajority() throws Exception {
        doTestOnPreparedThenOnCommittedWithNoPartitionId(MAJORITY);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionIdUnanimous() throws Exception {
        doTestOnPreparedThenOnCommittedWithNoPartitionId(UNANIMOUS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionMajority() {
        doTestOnPreparedThenOnCommittedWithNoPartition(MAJORITY);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoPartitionUnanimous() {
        doTestOnPreparedThenOnCommittedWithNoPartition(UNANIMOUS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoRemoteReplicasAlways() {
        doTestOnPreparedThenOnCommittedWithNoRemoteReplicas(ALWAYS);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoRemoteReplicasMajority() {
        doTestOnPreparedThenOnCommittedWithNoRemoteReplicas(MAJORITY);
    }

    @Test
    public void testOnPreparedThenOnCommittedWithNoRemoteReplicasUnanimous() {
        doTestOnPreparedThenOnCommittedWithNoRemoteReplicas(UNANIMOUS);
    }

    @Test
    public void testOnPreparedThenOnRolledBackWithNoPartitionIdMajority() throws Exception {
        doTestOnPreparedThenOnRolledBackWithNoPartitionId(MAJORITY);
    }

    @Test
    public void testOnPreparedThenOnRolledBackWithNoPartitionIdUnanimous() throws Exception {
        doTestOnPreparedThenOnRolledBackWithNoPartitionId(UNANIMOUS);
    }

    @Test
    public void testOnPreparedWithFailedReplicaRollsBackInUnanimousMode() throws Exception {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);

        Future<PrepareResult> future = createLeader(UNANIMOUS).onPrepared(repository, TX);
        assertFalse(future.get().isSuccess());

        verify(meshManager).getPartition(eq(PARTITION));
        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).rollBack(anyLong());
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(anyString());
    }

    @Test
    public void testOnPreparedWithFailedReplicaStillCommitsInAlwaysMode() throws Exception {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);

        LocalTransactionLeader leader = createLeader(ALWAYS);

        //Even though the replica fails to prepare the transaction, the local transaction should still
        //return success and commit
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertTrue(future.get().isSuccess());

        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock).asCompletionStage();
        verify(repositoryAttributeManager).incrementAndGetVersion(REPOSITORY_ID);
        verify(replicaTransaction, never()).commit(anyLong());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-1");
    }

    @Test
    public void testOnPreparedWithFailedReplicaStillCommitsInMajorityMode() throws Exception {
        List<ReplicaTransaction> replicaTransactions = setUpReplicaTransactions(FAILURE, SUCCESS);
        ReplicaTransaction replicaTransaction1 = replicaTransactions.get(0);
        ReplicaTransaction replicaTransaction2 = replicaTransactions.get(1);

        LocalTransactionLeader leader = createLeader(MAJORITY);

        //Even though replica1 fails to prepare the transaction, the local transaction should still
        //return success and commit because a majority votes YES (replica2 and local)
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertTrue(future.get().isSuccess());

        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction1).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction2).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock).asCompletionStage();
        verify(repositoryAttributeManager).incrementAndGetVersion(REPOSITORY_ID);
        verify(replicaTransaction1, never()).commit(anyLong());
        verify(replicaTransaction2).commit(anyLong());
        verify(replicaTransaction1).summarize(summaryCaptor.capture(), anyLong());
        verify(replicaTransaction2).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-1");
    }

    @Test
    public void testOnPreparedWithOnRolledBackCalled() {
        List<ReplicaTransaction> replicaTransactions = setUpReplicaTransactions(SUCCESS, SUCCESS);
        ReplicaTransaction replicaTransaction1 = replicaTransactions.get(0);
        ReplicaTransaction replicaTransaction2 = replicaTransactions.get(1);

        LocalTransactionLeader leader = createLeader(MAJORITY);
        leader.onRolledBack(null);

        Future<?> result = executorService.submit(() -> {
            try {
                Future<PrepareResult> prepareResultFuture = leader.onPrepared(repository, TX);
                assertFalse(prepareResultFuture.get().isSuccess());
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            } catch (ExecutionException e) {
                Throwables.propagateIfPossible(e.getCause(), RuntimeException.class);
                throw new RuntimeException(e.getCause());
            }
        });

        try {
            result.get();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (ExecutionException e) {
            Throwables.propagateIfPossible(e.getCause(), RuntimeException.class);
            throw new RuntimeException(e.getCause());
        }

        verify(replicaTransaction1).rollBack(anyLong());
        verify(replicaTransaction2).rollBack(anyLong());
    }

    @Test
    public void testOnPreparedWithSuccessfulReplicaAlways() throws Exception {
        doTestOnPreparedWithSuccessfulReplica(ALWAYS);
    }

    @Test
    public void testOnPreparedWithSuccessfulReplicaMajority() throws Exception {
        doTestOnPreparedWithSuccessfulReplica(MAJORITY);
    }

    @Test
    public void testOnPreparedWithSuccessfulReplicaUnanimous() throws Exception {
        doTestOnPreparedWithSuccessfulReplica(UNANIMOUS);
    }

    @Test
    public void testOnPreparedWithTwoFailedReplicaRollsBackInMajorityMode() throws Exception {
        List<ReplicaTransaction> replicaTransactions = setUpReplicaTransactions(FAILURE, FAILURE);
        ReplicaTransaction replicaTransaction1 = replicaTransactions.get(0);
        ReplicaTransaction replicaTransaction2 = replicaTransactions.get(1);

        LocalTransactionLeader leader = createLeader(MAJORITY);

        //When both replicas fail, the transaction should roll back
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertFalse(future.get().isSuccess());

        leader.onRolledBack(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction1).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock, never()).asCompletionStage();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(REPOSITORY_ID);
        verify(replicaTransaction1).rollBack(anyLong());
        verify(replicaTransaction2).rollBack(anyLong());
        verify(replicaTransaction1).summarize(summaryCaptor.capture(), anyLong());
        verify(replicaTransaction2).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertAllNodesConsistent(summary);
    }

    @Test
    public void testOnRolledBackHandlesRemoteFailuresMajority() {
        doTestOnRolledBackHandlesRemoteFailures(MAJORITY);
    }

    @Test
    public void testOnRolledBackHandlesRemoteFailuresUnanimous() {
        doTestOnRolledBackHandlesRemoteFailures(UNANIMOUS);
    }

    @Test
    public void testOnRolledBackMajority() {
        doTestOnRolledBack(MAJORITY);
    }

    @Test
    public void testOnRolledBackUnanimous() {
        doTestOnRolledBack(UNANIMOUS);
    }

    @Test
    public void testOnRolledBackWhenLocallyInconsistentMajority() {
        doTestOnRolledBackWhenLocallyInconsistent(MAJORITY);
    }

    @Test
    public void testOnRolledBackWhenLocallyInconsistentUnanimous() {
        doTestOnRolledBackWhenLocallyInconsistent(UNANIMOUS);
    }

    @Test
    public void testOnRolledBackWhenRemotelyInconsistentMajority() {
        doTestOnRolledBackWhenRemotelyInconsistent(MAJORITY);
    }

    @Test
    public void testOnRolledBackWhenRemotelyInconsistentUnanimous() {
        doTestOnRolledBackWhenRemotelyInconsistent(UNANIMOUS);
    }

    @Test
    public void testTriggersContentHashRecalculationWhenWriteLockHeld() {
        writeLockFuture = new CompletableFuture<>();
        setUpReplicaTransactions(SUCCESS);
        LocalTransactionLeader leader = createLeader(MAJORITY);

        Future<PrepareResult> future = leader.onPrepared(repository, TX);

        //Verify the transaction is prepared but write lock was not acquired yet
        assertFalse(future.isDone());
        verify(writeLock).asCompletionStage();
        verify(repositoryHashService, never()).maybeRecalculateContentHash(repository);

        //Write lock was acquired
        writeLockFuture.complete(null);

        //Now we should have recalculated the content hash
        verify(repositoryHashService).maybeRecalculateContentHash(repository);
    }

    @Test
    public void testWithLocallyInconsistentVersionAlwaysMode() {
        doTestWithLocallyInconsistentVersion(ALWAYS);
    }

    @Test
    public void testWithLocallyInconsistentVersionMajorityMode() {
        doTestWithLocallyInconsistentVersion(MAJORITY);
    }

    @Test
    public void testWithLocallyInconsistentVersionUnanimousMode() {
        doTestWithLocallyInconsistentVersion(MAJORITY);
    }

    private static MeshNode mockNode(int id) {
        MeshNode node = mock(MeshNode.class);
        when(node.getId()).thenReturn("remote-" + id);
        when(node.isRemote()).thenReturn(true);
        return node;
    }

    private static PartitionReplica mockReplica(MeshNode node) {
        PartitionReplica replica = mock(PartitionReplica.class);
        when(replica.getNode()).thenReturn(node);
        return replica;
    }

    private void assertAllNodesConsistent(RpcTransactionSummary summary) {
        assertThat(summary.getInconsistentReplicasList(), empty());
        verifyNoInteractions(sidebandManager);
    }

    private void assertMarkedInconsistent(RpcTransactionSummary summary, String... nodeIds) {
        if (nodeIds.length == 0) {
            assertAllNodesConsistent(summary);
        } else {
            assertThat(summary.getInconsistentReplicasList(), containsInAnyOrder(nodeIds));
            Set<String> expectedNodeIds = ImmutableSet.copyOf(nodeIds);
            ArgumentCaptor<RpcSidebandFollowerFragment> fragmentCaptor =
                    ArgumentCaptor.forClass(RpcSidebandFollowerFragment.class);
            verify(sidebandManager, atLeastOnce()).send(fragmentCaptor.capture());
            Set<String> reportedNodes = fragmentCaptor.getAllValues().stream()
                    .map(RpcSidebandFollowerFragment::getRepositoryReplicaObservation)
                    .filter(observation -> observation.getState() != RpcReplicaState.REPLICA_STATE_CONSISTENT)
                    .flatMap(observation -> observation.getNodesList().stream())
                    .collect(Collectors.toSet());
            assertEquals(expectedNodeIds, reportedNodes);
        }
    }

    private LocalTransactionLeader createLeader(QuorumMode mode) {
        when(replicationConfig.getQuorumMode()).thenReturn(mode);

        return new LocalTransactionLeader("tx", config, eventPublisher, logicalClock, meshClient, meshManager,
                repositoryAttributeManager, repositoryHashService, sidebandManager, writeTokenRegistry);
    }

    private void doTestOnCommitted(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        InOrder inOrder = inOrder(replicaTransaction, writeLock, repositoryAttributeManager);
        inOrder.verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        inOrder.verify(writeLock).asCompletionStage();
        inOrder.verify(replicaTransaction).commit(anyLong());
        inOrder.verify(repositoryAttributeManager).incrementAndGetVersion(REPOSITORY_ID);
        inOrder.verify(replicaTransaction).await(notNull());
        inOrder.verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        inOrder.verify(writeLock).close();

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertAllNodesConsistent(summary);
    }

    private void doTestOnCommittedHandlesRemoteFailures(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);
        when(replicaTransaction.await(notNull())).thenThrow(TransactionCommitFailedException.class);
        when(replicaTransaction.getReplicaVersion()).thenReturn(Optional.empty());

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).commit(anyLong());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-1");

        logListener.assertWarn("Failed to commit");
    }

    private void doTestOnCommittedWhenRemotelyInconsistent(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);
        when(replicaTransaction.await(notNull())).thenReturn(false);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).commit(anyLong());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-1");
    }

    private void doTestOnCommittedWhenRemotelyInconsistentVersion(QuorumMode quorumMode) {
        List<ReplicaTransaction> replicaTransactions =
                setUpReplicaTransactions(SUCCESS, SUCCESS);
        ReplicaTransaction replicaTransaction1 = replicaTransactions.get(0);
        ReplicaTransaction replicaTransaction2 = replicaTransactions.get(1);

        //Mock the one of the remote replicas returning an older version
        Optional<ReplicaVersion> staleVersion = of(new ReplicaVersion(CONTENT_HASH, VERSION));
        doReturn(staleVersion).when(replicaTransaction2).getReplicaVersion();

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction1).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction1).commit(anyLong());
        verify(replicaTransaction1).await(notNull());
        verify(replicaTransaction1).summarize(summaryCaptor.capture(), anyLong());

        verify(replicaTransaction2).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction2).commit(anyLong());
        verify(replicaTransaction2).await(notNull());
        verify(replicaTransaction2).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-2");
    }

    //This test sets up 1 follower, which abstains from the vote. Abstaining nodes should be removed from both
    //the write lock _and the replica count_, which means the replica count should end up 1, instead of 2, and
    //MAJORITY and UNANIMOUS modes should succeed (along with ALWAYS) where, had the node voted NO instead, both
    //would end up rolling back (since a replica count of 2 requires both replicas to vote YES for those modes)
    private void doTestOnCommittedWhenReplicaAbstains(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(ABSTAIN);
        when(replicaTransaction.await(notNull())).thenReturn(false);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).rollBack(anyLong());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "remote-1");
    }

    private void doTestOnCommittedWithLocalFailure(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(new RuntimeException());

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction).commit(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(true));
        assertThat(summary.getContentHash(), equalTo(CONTENT_HASH));
        assertMarkedInconsistent(summary, "local");

        ArgumentCaptor<ReplicaStateObservedEvent> eventCaptor =
                ArgumentCaptor.forClass(ReplicaStateObservedEvent.class);
        verify(eventPublisher).publish(eventCaptor.capture());

        assertEquals(ReplicaState.INCONSISTENT, eventCaptor.getValue().getObservation().getState());
    }

    private void doTestOnPreparedTakesPartitionMigrationsIntoAccountAllQuorumsPass(QuorumMode quorumMode)
            throws ExecutionException, InterruptedException {
        //Set up 3 replicas and a partition migration from 2 -> 3, and have both 2 and 3 fail to replicate
        //In majority mode, this should still pass because
        List<ReplicaTransaction> transactions = setUpReplicaTransactions(SUCCESS, FAILURE, FAILURE);
        mockManagedRepository();
        MeshNode node2 = transactions.get(1).getNode();
        MeshNode node3 = transactions.get(2).getNode();
        PartitionReplicaMigration migration = new PartitionReplicaMigration.Builder()
                .partition(PARTITION)
                .sourceNode(node2)
                .targetNode(node3)
                .build();
        when(meshManager.getPartitionMigrations("partition")).thenReturn(ImmutableList.of(migration));

        LocalTransactionLeader leader = createLeader(quorumMode);
        Future<PrepareResult> future = leader.onPrepared(repository, TX);

        //verify that each replica transaction is instructed to commit/rollback
        boolean commitExpected = quorumMode != UNANIMOUS;
        assertEquals(commitExpected, future.get().isSuccess());
        for (ReplicaTransaction transaction : transactions) {
            if (commitExpected && !node2.equals(transaction.getNode()) && !node3.equals(transaction.getNode())) {
                verify(transaction).commit(anyLong());
            } else {
                verify(transaction).rollBack(anyLong());
            }
        }
    }

    private void doTestOnPreparedThenOnCommittedWithNoPartition(QuorumMode quorumMode) {
        when(repository.accept(any())).thenReturn("partition");

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        verify(meshManager).getPartition(eq("partition"));
        verifyNoInteractions(meshClient, writeTokenRegistry);
    }

    private void doTestOnPreparedThenOnCommittedWithNoPartitionId(QuorumMode quorumMode) throws Exception {
        LocalTransactionLeader leader = createLeader(quorumMode);
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertTrue(future.get().isSuccess());

        leader.onCommitted(null);

        verify(meshManager).getLocalNode();
        verifyNoMoreInteractions(meshClient, meshManager, writeTokenRegistry);
    }

    private void doTestOnPreparedThenOnCommittedWithNoRemoteReplicas(QuorumMode quorumMode) {
        setUpMockPartition(LOCAL);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onCommitted(null);

        verify(meshManager, atLeastOnce()).getLocalNode();
        verify(writeLock).asCompletionStage();
        verify(repositoryAttributeManager).incrementAndGetVersion(REPOSITORY_ID);
        verify(writeLock).close();
        verifyNoInteractions(meshClient);
    }

    private void doTestOnPreparedThenOnRolledBackWithNoPartitionId(QuorumMode quorumMode) throws Exception {
        LocalTransactionLeader leader = createLeader(quorumMode);
        Future<PrepareResult> future = leader.onPrepared(repository, TX);
        assertTrue(future.get().isSuccess());

        leader.onRolledBack(null);

        verify(meshManager).getLocalNode();
        verifyNoMoreInteractions(meshClient, meshManager, writeTokenRegistry);
    }

    private void doTestOnPreparedWithSuccessfulReplica(QuorumMode quorumMode) throws Exception {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(SUCCESS);

        Future<PrepareResult> future = createLeader(quorumMode).onPrepared(repository, TX);
        assertTrue(future.get().isSuccess());

        InOrder inOrder = inOrder(meshManager, replicaTransaction, writeLock, repositoryAttributeManager);
        inOrder.verify(meshManager).getPartition(eq("partition"));
        inOrder.verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        inOrder.verify(writeLock).asCompletionStage();
        inOrder.verify(replicaTransaction).commit(anyLong());
    }

    private void doTestOnRolledBack(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onRolledBack(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock, never()).asCompletionStage();
        verify(replicaTransaction).rollBack(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(anyString());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertAllNodesConsistent(summary);
    }

    private void doTestOnRolledBackHandlesRemoteFailures(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);
        when(replicaTransaction.await(notNull())).thenThrow(TransactionRollbackFailedException.class);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onRolledBack(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock, never()).asCompletionStage();
        verify(replicaTransaction).rollBack(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(anyString());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertAllNodesConsistent(summary);

        logListener.assertWarn("Failed to roll back");
    }

    private void doTestOnRolledBackWhenLocallyInconsistent(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onRolledBack(new RuntimeException());

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock, never()).asCompletionStage();
        verify(replicaTransaction).rollBack(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(anyString());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertAllNodesConsistent(summary);
    }

    private void doTestOnRolledBackWhenRemotelyInconsistent(QuorumMode quorumMode) {
        ReplicaTransaction replicaTransaction = setUpReplicaTransaction(FAILURE);
        when(replicaTransaction.await(notNull())).thenReturn(true);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        leader.onRolledBack(null);

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(writeLock, never()).asCompletionStage();
        verify(replicaTransaction).rollBack(anyLong());
        verify(replicaTransaction).await(notNull());
        verify(replicaTransaction).summarize(summaryCaptor.capture(), anyLong());
        verify(writeLock).close();
        verify(repositoryAttributeManager, never()).incrementAndGetVersion(anyString());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(false));
        assertThat(summary.getContentHash(), equalTo(""));
        assertMarkedInconsistent(summary, "remote-1");
    }

    //This test case doesn't call out onCommitted or onRolledBack specifically because which gets executed
    //depends on the quorum mode
    private void doTestWithLocallyInconsistentVersion(QuorumMode quorumMode) {
        List<ReplicaTransaction> replicaTransactions = setUpReplicaTransactions(SUCCESS, SUCCESS);
        ReplicaTransaction replicaTransaction1 = replicaTransactions.get(0);
        ReplicaTransaction replicaTransaction2 = replicaTransactions.get(1);

        //Mock the local replica being at an older version than the remote replicas
        when(repositoryAttributeManager.getVersion(REPOSITORY_ID)).thenReturn(VERSION - 1L);
        when(repositoryAttributeManager.incrementAndGetVersion(REPOSITORY_ID)).thenReturn(VERSION);

        LocalTransactionLeader leader = createLeader(quorumMode);
        leader.onPrepared(repository, TX);
        //In ALWAYS mode, if the (inconsistent) leader can commit, we will commit.
        //In MAJORITY or UNANIMOUS mode, we roll back if the leader cannot commit.
        boolean commitExpected = quorumMode == ALWAYS;
        if (commitExpected) {
            leader.onCommitted(null);
        } else {
            leader.onRolledBack(null);
        }

        ArgumentCaptor<RpcTransactionSummary> summaryCaptor = ArgumentCaptor.forClass(RpcTransactionSummary.class);

        verify(replicaTransaction1).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction1).rollBack(anyLong());
        verify(replicaTransaction1).await(notNull());
        verify(replicaTransaction1).summarize(summaryCaptor.capture(), anyLong());

        verify(replicaTransaction2).prepare(eq(TX_FRAGMENTS), eq(TOKEN_TIMESTAMP));
        verify(replicaTransaction2).rollBack(anyLong());
        verify(replicaTransaction2).await(notNull());
        verify(replicaTransaction2).summarize(summaryCaptor.capture(), anyLong());

        RpcTransactionSummary summary = summaryCaptor.getValue();
        assertThat(summary.getCommitted(), equalTo(commitExpected));
        assertThat(summary.getContentHash(), equalTo(commitExpected ? CONTENT_HASH : ""));
        if (commitExpected) {
            //In ALWAYS mode, the leader dictates what gets committed and therefore which version
            //is considered to be the consistent version
            assertMarkedInconsistent(summary, "remote-1", "remote-2");
        } else {
            //In other modes, since we roll back, every node should roll back successfully and
            //the leader (local) should report as inconsistent
            assertMarkedInconsistent(summary, "local");
        }
    }

    private void mockManagedRepository() {
        ManagedRepository repository = mock(ManagedRepository.class);
        when(repository.accept(any())).thenReturn(PARTITION);
        when(repository.getId()).thenReturn(REPOSITORY_ID);
        when(repository.getPartition()).thenReturn(PARTITION);
        this.repository = repository;
    }

    private MockReplicaTransaction mockReplicaTransaction(int id, MeshNode node, PrepareResult result,
                                                          boolean hasWriteToken) {
        ReplicaTransaction replicaTransaction = mock(ReplicaTransaction.class, "remote-tx-" + id);
        when(replicaTransaction.getNode()).thenReturn(node);

        MutableObject<ReplicaTransactionCallback> callbackRef = new MutableObject<>();
        MutableBoolean prepared = new MutableBoolean();
        when(replicaTransaction.prepare(anyCollection(), anyLong())).thenAnswer(ignored -> {
            if (result.isSuccess()) {
                ReplicaTransactionCallback callback = callbackRef.getValue();
                if (hasWriteToken) {
                    callback.onTokenGranted(node, PREPARED_REPLICA_VERSION);
                }
            }
            prepared.setTrue();

            return CompletableFuture.completedFuture(result);
        });
        when(replicaTransaction.getPrepareResult()).thenAnswer(ignored ->
                prepared.booleanValue() ? of(result) : Optional.empty());

        MutableBoolean committed = new MutableBoolean();
        doAnswer(ignored -> {
            committed.setTrue();
            return null;
        }).when(replicaTransaction).commit(anyLong());
        when(replicaTransaction.getReplicaVersion()).thenAnswer(ignored ->
                committed.booleanValue() ? of(new ReplicaVersion(CONTENT_HASH, VERSION + 1)) : Optional.empty());
        when(replicaTransaction.await(notNull())).thenAnswer(ignored -> committed.booleanValue());

        when(meshClient.replicate(same(node), any())).thenAnswer(invocation -> {
            ReplicateParameters parameters = invocation.getArgument(1);
            callbackRef.setValue(parameters.getCallback());
            return replicaTransaction;
        });
        return new MockReplicaTransaction() {

            @Override
            public ReplicaTransactionCallback getCallback() {
                return callbackRef.getValue();
            }

            @Override
            public ReplicaTransaction getMock() {
                return replicaTransaction;
            }

            @Override
            public void grantWriteToken() {
                if (callbackRef.getValue() == null) {
                    throw new IllegalStateException("callback is null");
                }
                callbackRef.getValue().onTokenGranted(node, PREPARED_REPLICA_VERSION);
            }
        };
    }

    private void setUpMockPartition(MeshNode node) {
        setUpMockPartition(Collections.singleton(node));
    }

    private void setUpMockPartition(Collection<MeshNode> nodes) {
        List<PartitionReplica> replicas = new ArrayList<>();
        replicas.add(mockReplica(LOCAL));
        for (MeshNode node : nodes) {
            replicas.add(mockReplica(node));
        }

        Partition partition = mock(Partition.class);
        when(meshManager.getPartition(eq("partition"))).thenReturn(partition);
        when(partition.getReplicas()).thenReturn(replicas);
        when(repository.accept(any())).thenReturn("partition");
    }

    private ReplicaTransaction setUpReplicaTransaction(PrepareResult result) {
        return setUpReplicaTransactions(result).get(0);
    }

    private List<MockReplicaTransaction> setUpReplicaTransactions(List<Pair<PrepareResult, Boolean>> results) {
        List<MockReplicaTransaction> transactions = new ArrayList<>();
        List<MeshNode> nodes = new ArrayList<>();
        int counter = 1;

        for (Pair<PrepareResult, Boolean> result : results) {
            int id = counter++;
            MeshNode node = mockNode(id);
            MockReplicaTransaction tx = mockReplicaTransaction(id, node, result.getLeft(), result.getRight());
            nodes.add(node);
            transactions.add(tx);
        }

        setUpMockPartition(nodes);
        return transactions;
    }

    private List<ReplicaTransaction> setUpReplicaTransactions(PrepareResult... results) {
        List<ReplicaTransaction> transactions = new ArrayList<>();
        List<MeshNode> nodes = new ArrayList<>();
        int counter = 1;

        for (PrepareResult result : results) {
            int id = counter++;
            MeshNode node = mockNode(id);
            ReplicaTransaction tx = mockReplicaTransaction(id, node, result, true).getMock();
            nodes.add(node);
            transactions.add(tx);
        }

        setUpMockPartition(nodes);
        return transactions;
    }

    interface MockReplicaTransaction {

        ReplicaTransactionCallback getCallback();

        ReplicaTransaction getMock();

        void grantWriteToken();
    }
}