package com.palantir.atlasdb.keyvalue.impl.partition;

import com.google.common.collect.ImmutableSet;
import com.palantir.atlasdb.keyvalue.api.KeyValueService;
import com.palantir.atlasdb.keyvalue.impl.AbstractAtlasDbKeyValueServiceTest;
import com.palantir.atlasdb.keyvalue.impl.InMemoryKeyValueService;
import com.palantir.atlasdb.keyvalue.partition.FailableKeyValueService;
import com.palantir.atlasdb.keyvalue.partition.PartitionedKeyValueService;

public class AvailabilityTest extends AbstractAtlasDbKeyValueServiceTest {

    KeyValueService kvs;
    public static final int NUM_SERVICES = 4;
    static final FailableKeyValueService[] services = new FailableKeyValueService[4];

    int numBroken = 0;

    @Override
    protected KeyValueService getKeyValueService() {
        if (kvs == null) {
            createServices();
        }
        return kvs;
    }

    private void createServices() {
        for (int i = 0; i < NUM_SERVICES; ++i) {
            FailableKeyValueService service = FailableKeyValueService.wrap(new InMemoryKeyValueService(
                    false));
            services[i] = service;
        }
//        services[0].stop();
//        services[1].stop();
        services[2].stop();
//        services[3].stop();
        kvs = PartitionedKeyValueService.create(ImmutableSet.<FailableKeyValueService> copyOf(services));
    }

}

@@ -59,46 +59,133 @@
 ort 
-com.google.common.collect.ImmutableSet
+java.util.Random;%0Aimport java.util.Set;%0A%0Aimport com.google.common.collect.ImmutableSet;%0Aimport com.google.common.collect.Sets
 ;%0Aim
@@ -531,16 +531,81 @@
 Service;
+%0Aimport com.palantir.atlasdb.keyvalue.partition.QuorumParameters;
 %0A%0Apublic
@@ -742,17 +742,17 @@
 VICES = 
-4
+5
 ;%0A    st
@@ -831,34 +831,165 @@
 ice%5B
-4%5D;%0A%0A    int numBroken = 0
+NUM_SERVICES%5D;%0A    QuorumParameters quorumParameters = new QuorumParameters(5, 3, 3);%0A    // TODO: Make it deterministic (?)%0A    Random random = new Random()
 ;%0A%0A 
@@ -1148,16 +1148,557 @@
 %0A    %7D%0A%0A
+    private Set%3CInteger%3E generateBrokenServices() %7B%0A        return generateBrokenServices(%0A                Sets.%3CInteger%3E newHashSet(),%0A                quorumParameters.getReplicationFactor() - quorumParameters.getReadFactor());%0A    %7D%0A%0A    private Set%3CInteger%3E generateBrokenServices(Set%3CInteger%3E brokenReads, int numBreaks) %7B%0A        if (brokenReads.size() == numBreaks) %7B%0A            return brokenReads;%0A        %7D%0A        brokenReads.add(random.nextInt(NUM_SERVICES));%0A        return generateBrokenServices(brokenReads, numBreaks);%0A    %7D%0A%0A
     priv
@@ -1952,26 +1952,24 @@
       %7D%0A
-//
         
 services
@@ -1964,38 +1964,82 @@
     
-services%5B0%5D.stop
+Set%3C? extends Integer%3E brokenServices = generateBrokenServices
 ();%0A
-//
         
 serv
@@ -2038,28 +2038,47 @@
     
-services%5B1%5D.stop();%0A
+for (Integer i : brokenServices) %7B%0A    
     
@@ -2094,21 +2094,36 @@
 ces%5B
-2%5D.stop();%0A//
+i%5D.setBrokenRead(true);%0A    
     
@@ -2139,18 +2139,42 @@
 ces%5B
-3%5D.stop();
+i%5D.setBrokenWrite(true);%0A        %7D
 %0A   
@@ -2218,16 +2218,33 @@
 .create(
+%0A                
 Immutabl
@@ -2290,16 +2290,63 @@
 ervices)
+,%0A                new QuorumParameters(5, 3, 3)
 );%0A    %7D

Fix Availability test.