package com.palantir.atlasdb.keyvalue.remoting;

import java.lang.reflect.Field;

import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.datatype.guava.GuavaModule;
import com.google.common.base.Preconditions;
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.KeyValueEndpoint;
import com.palantir.atlasdb.keyvalue.partition.PartitionMapService;
import com.palantir.atlasdb.keyvalue.partition.PartitionMapServiceImpl;
import com.palantir.atlasdb.keyvalue.partition.QuorumParameters;
import com.palantir.atlasdb.keyvalue.partition.SimpleKeyValueEndpoint;

import io.dropwizard.Configuration;
import io.dropwizard.testing.DropwizardTestSupport;
import io.dropwizard.testing.junit.DropwizardClientRule;

public class KeyValueEndpointTest extends AbstractAtlasDbKeyValueServiceTest {

    static final KeyValueService endpointKvs = RemotingKeyValueService.createServerSide(new InMemoryKeyValueService(false), null);
    static final PartitionMapService endpointPms = new PartitionMapServiceImpl();
    static final QuorumParameters QUORUM_PARAMETERS = new QuorumParameters(3, 2, 2);

    private final SimpleModule module = RemotingKeyValueService.kvsModule();
    private final ObjectMapper mapper = RemotingKeyValueService.kvsMapper();

    @Rule public final DropwizardClientRule endpointKvsService = new DropwizardClientRule(endpointKvs);
    @Rule public final DropwizardClientRule endpointPmsService = new DropwizardClientRule(endpointPms);

    KeyValueEndpoint endpoint;

    @SuppressWarnings("unchecked")
    @Before
    public void setupHacks() throws Exception {
        Field field = endpointKvsService.getClass().getDeclaredField("testSupport");
        field.setAccessible(true);
        DropwizardTestSupport<Configuration> testSupport = (DropwizardTestSupport<Configuration>) field.get(endpointKvsService);
        ObjectMapper mapper = testSupport.getEnvironment().getObjectMapper();
        mapper.registerModule(module);
        mapper.registerModule(new GuavaModule());
        testSupport.getApplication();
    }

    @Test
    public void testSimple() {
    }

    @Override
    protected KeyValueService getKeyValueService() {
        if (endpoint == null) {
            endpoint = new SimpleKeyValueEndpoint(endpointKvsService.baseUri().toString(), endpointPmsService.baseUri().toString());
        }
        return Preconditions.checkNotNull(endpoint.keyValueService());
    }

}

@@ -42,16 +42,62 @@
 oting;%0A%0A
+import static org.junit.Assert.assertEquals;%0A%0A
 import j
@@ -2436,79 +2436,15 @@
     
-@Test%0A    public void testSimple() %7B%0A    %7D%0A%0A    @Override%0A    protec
+priva
 te
-d
  Key
@@ -2452,34 +2452,28 @@
 alue
-Service getKeyValueService
+Endpoint getEndpoint
 () %7B
@@ -2648,16 +2648,239 @@
       %7D%0A
+        return endpoint;%0A    %7D%0A%0A    @Test%0A    public void testSimple() %7B%0A        assertEquals(0L, getEndpoint().partitionMapService().getVersion());%0A    %7D%0A%0A    @Override%0A    protected KeyValueService getKeyValueService() %7B%0A
         
@@ -2913,24 +2913,29 @@
 NotNull(
-e
+getE
 ndpoint
+()
 .keyValu

Update KVE test.