Fix consul value fetch

cfg4j-consul/src/main/java/org/cfg4j/source/consul/ConsulConfigurationSource.java

@@ -17,6 +17,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import com.google.common.net.HostAndPort;
 import com.orbitz.consul.Consul;
 import com.orbitz.consul.KeyValueClient;
 import com.orbitz.consul.model.kv.Value;
@@ -33,8 +34,16 @@
 
 /**
  * Note: use {@link ConsulConfigurationSourceBuilder} for building instances of this class.
+ *
+ * <p>
+ * This source interpretes the provided {@link Environment} variable as a Consul Key Prefix.
+ * If you build this Source without passing the {@link Environment} to {@link ConsulConfigurationSourceBuilder#withEnvironment(Environment)}
+ * then by default all values will be fetched from Consul to be matched.
+ * </p>
+ *
  * <p>
  * Read configuration from the Consul K-V store.
+ * </p>
  */
 class ConsulConfigurationSource implements ConfigurationSource {
 
@@ -45,6 +54,7 @@ class ConsulConfigurationSource implements ConfigurationSource {
   private final String host;
   private final int port;
   private boolean initialized;
+  private final Environment environment;
 
   /**
    * Note: use {@link ConsulConfigurationSourceBuilder} for building instances of this class.
@@ -54,9 +64,10 @@ class ConsulConfigurationSource implements ConfigurationSource {
    * @param host Consul host to connect to
    * @param port Consul port to connect to
    */
-  ConsulConfigurationSource(String host, int port) {
+  ConsulConfigurationSource(String host, int port, Environment environment) {
     this.host = requireNonNull(host);
     this.port = port;
+    this.environment = environment;
 
     initialized = false;
   }
@@ -70,19 +81,11 @@ public Properties getConfiguration(Environment environment) {
     }
 
     Properties properties = new Properties();
-    String path = environment.getName();
-
-    if (path.startsWith("/")) {
-      path = path.substring(1);
-    }
-
-    if (path.length() > 0 && !path.endsWith("/")) {
-      path = path + "/";
-    }
+    String environmentPrefix = formatEnvironment(environment);
 
     for (Map.Entry<String, String> entry : consulValues.entrySet()) {
-      if (entry.getKey().startsWith(path)) {
-        properties.put(entry.getKey().substring(path.length()).replace("/", "."), entry.getValue());
+      if (entry.getKey().startsWith(environmentPrefix)) {
+        properties.put(entry.getKey().substring(environmentPrefix.length()).replace("/", "."), entry.getValue());
       }
     }
 
@@ -97,7 +100,8 @@ public void init() {
     try {
       LOG.info("Connecting to Consul client at " + host + ":" + port);
 
-      Consul consul = Consul.newClient(host, port);
+      Consul consul =  Consul.builder().withHostAndPort(HostAndPort.fromParts(host,port)).build();
+
       kvClient = consul.keyValueClient();
     } catch (Exception e) {
       throw new SourceCommunicationException("Can't connect to host " + host + ":" + port, e);
@@ -113,8 +117,9 @@ public void reload() {
     List<Value> valueList;
 
     try {
-      LOG.debug("Reloading configuration from Consuls' K-V store");
-      valueList = kvClient.getValues("/");
+      String consulKeyPath = formatEnvironmentToConsulKeyPrefix(environment);
+      LOG.debug("Reloading configuration from Consuls' K-V store for Prefix: " + consulKeyPath);
+      valueList = kvClient.getValues(consulKeyPath);
     } catch (Exception e) {
       initialized = false;
       throw new SourceCommunicationException("Can't get values from k-v store", e);
@@ -142,4 +147,31 @@ public String toString() {
         ", kvClient=" + kvClient +
         '}';
   }
+
+  private static String formatEnvironmentToConsulKeyPrefix(Environment environment) {
+    String path = environment.getName();
+
+    if (path.startsWith("/")) {
+      path = path.substring(1);
+    }
+
+    if(path.length() > 0 && path.endsWith("/")) {
+      path = path.substring(0, path.length() - 1);
+    }
+    return path.isEmpty() ? "/" : path;
+  }
+
+  private static String formatEnvironment(Environment environment) {
+    String path = environment.getName();
+
+    if (path.startsWith("/")) {
+      path = path.substring(1);
+    }
+
+    if (path.length() > 0 && !path.endsWith("/")) {
+      path = path + "/";
+    }
+
+    return path;
+  }
 }

cfg4j-consul/src/main/java/org/cfg4j/source/consul/ConsulConfigurationSourceBuilder.java

@@ -15,13 +15,17 @@
  */
 package org.cfg4j.source.consul;
 
+import org.cfg4j.source.context.environment.Environment;
+import org.cfg4j.source.context.environment.ImmutableEnvironment;
+
 /**
  * Builder for {@link ConsulConfigurationSource}.
  */
 public class ConsulConfigurationSourceBuilder {
 
   private String host;
   private int port;
+  private Environment environment;
 
   /**
    * Construct {@link ConsulConfigurationSource}s builder
@@ -35,6 +39,7 @@ public class ConsulConfigurationSourceBuilder {
   public ConsulConfigurationSourceBuilder() {
     host = "localhost";
     port = 8500;
+    environment = new ImmutableEnvironment("/");
   }
 
   /**
@@ -59,20 +64,32 @@ public ConsulConfigurationSourceBuilder withPort(int port) {
     return this;
   }
 
+  /**
+   * Set the {@link Environment} to use as your Consul Key Prefix root for configuration.
+   *
+   * @param environment a {@link Environment} interpreted as a Consul Key Prefix
+   * @return this builder with the Environment set to the provided Environment
+     */
+  public ConsulConfigurationSourceBuilder withEnvironment(Environment environment) {
+    this.environment = environment;
+    return this;
+  }
+
   /**
    * Build a {@link ConsulConfigurationSource} using this builder's configuration
    *
    * @return new {@link ConsulConfigurationSource}
    */
   public ConsulConfigurationSource build() {
-    return new ConsulConfigurationSource(host, port);
+    return new ConsulConfigurationSource(host, port,environment);
   }
 
   @Override
   public String toString() {
     return "ConsulConfigurationSource{" +
         "host=" + host +
         ", port=" + port +
+        ", environment=" + environment +
         '}';
   }
 }

cfg4j-consul/src/test/java/org/cfg4j/source/consul/ConsulConfigurationSourceIntegrationTest.java

@@ -31,6 +31,7 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
+import org.mockito.internal.exceptions.ExceptionIncludingMockitoWarnings;
 import org.mockito.runners.MockitoJUnitRunner;
 
 import java.io.IOException;
@@ -50,6 +51,11 @@ private class ModifiableDispatcher extends Dispatcher {
     private static final String disabledBase64 = "ZGlzYWJsZWQ=";
     private static final String enabledBase64 = "ZW5hYmxlZA==";
 
+    private static final String usWest1Config = "{\"CreateIndex\":1,\"ModifyIndex\":1,\"LockIndex\":0,\"Key\":\"us-west-1/featureA.toggle\",\"Flags\":0,\"Value\":\"" + disabledBase64 + "\"}";
+    private static final String usWest2FeatureEnable = "{\"CreateIndex\":2,\"ModifyIndex\":2,\"LockIndex\":0,\"Key\":\"us-west-2/featureB.toggle\",\"Flags\":0,\"Value\":\"" + enabledBase64 + "\"}";
+    private static final String usWest2FeatureDisabled = "{\"CreateIndex\":2,\"ModifyIndex\":2,\"LockIndex\":0,\"Key\":\"us-west-2/featureB.toggle\",\"Flags\":0,\"Value\":\"" + disabledBase64 + "\"}";
+
+
     private boolean usWest2Toggle = false;
 
     void toggleUsWest2() {
@@ -58,17 +64,30 @@ void toggleUsWest2() {
 
     @Override
     public MockResponse dispatch(RecordedRequest request) throws InterruptedException {
-
+      StringBuilder sbResponseBody = new StringBuilder();
       switch (request.getPath()) {
         case "/v1/agent/self":
           return new MockResponse().setResponseCode(200).setBody(PING_RESPONSE);
         case "/v1/kv/?recurse=true":
+          sbResponseBody.append("[");
+          sbResponseBody.append(usWest1Config).append(",");
+          if(usWest2Toggle) {
+            sbResponseBody.append(usWest2FeatureEnable);
+          } else {
+            sbResponseBody.append(usWest2FeatureDisabled);
+          }
+          sbResponseBody.append("]");
+
           return new MockResponse()
               .setResponseCode(200)
               .addHeader("Content-Type", "application/json; charset=utf-8")
-              .setBody("[{\"CreateIndex\":1,\"ModifyIndex\":1,\"LockIndex\":0,\"Key\":\"us-west-1/featureA.toggle\",\"Flags\":0,\"Value\":\"ZGlzYWJsZWQ=\"},"
-                  + "{\"CreateIndex\":2,\"ModifyIndex\":2,\"LockIndex\":0,\"Key\":\"us-west-2/featureA.toggle\",\"Flags\":0,\"Value\":\""
-                  + (usWest2Toggle ? enabledBase64 : disabledBase64) + "\"}]");
+              .setBody(sbResponseBody.toString());
+
+        case "/v1/kv/us-west-1?recurse=true":
+          return new MockResponse()
+                  .setResponseCode(200)
+                  .addHeader("Content-Type", "application/json; charset=utf-8")
+                  .setBody("[" + usWest1Config + "]");
       }
       return new MockResponse().setResponseCode(404);
     }
@@ -130,14 +149,33 @@ public void getConfigurationShouldIgnoreLeadingSlashInGivenEnvironment() throws
     assertThat(source.getConfiguration(environment)).contains(MapEntry.entry("featureA.toggle", "disabled"));
   }
 
+  @Test
+  public void getConfigurationWithSourceEnvironmentSetShouldReturnOnlyKeysInEnvironment() throws Exception {
+    Environment environment = new ImmutableEnvironment("/us-west-1");
+    Environment noEnvironment = new ImmutableEnvironment("");
+
+    ConsulConfigurationSource source = new ConsulConfigurationSourceBuilder()
+            .withHost(server.getHostName())
+            .withPort(server.getPort())
+            .withEnvironment(environment)
+            .build();
+
+    source.init();
+
+    // Match with any prefix, we shouldn't have gotten us-west-2 data back
+    assertThat(source.getConfiguration(noEnvironment)).doesNotContain(MapEntry.entry("featureB.toggle", "disabled"));
+
+    assertThat(source.getConfiguration(environment)).contains(MapEntry.entry("featureA.toggle", "disabled"));
+  }
+
   @Test
   public void getConfigurationShouldBeUpdatedByReload() throws Exception {
     dispatcher.toggleUsWest2();
 
     source.reload();
 
     Environment environment = new ImmutableEnvironment("us-west-2");
-    assertThat(source.getConfiguration(environment)).contains(MapEntry.entry("featureA.toggle", "enabled"));
+    assertThat(source.getConfiguration(environment)).contains(MapEntry.entry("featureB.toggle", "enabled"));
   }
 
   @Test