From 18456bccbf6fda41233c05802845a0671f7e5c41 Mon Sep 17 00:00:00 2001
From: "Fitzgerald, Andrew" <andrewcfitzgerald@gmail.com>
Date: Sat, 2 Apr 2016 01:09:42 -0400
Subject: [PATCH] Adds configurable default service tag for queries.

---
 .../discovery/ConsulDiscoveryProperties.java  | 15 ++++++
 .../consul/discovery/ConsulServerList.java    |  2 +-
 .../ConsulDiscoveryPropertiesTests.java       | 46 +++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryPropertiesTests.java

diff --git a/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryProperties.java b/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryProperties.java
index c8df745b..a3bb7c02 100644
--- a/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryProperties.java
+++ b/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryProperties.java
@@ -114,6 +114,11 @@ public class ConsulDiscoveryProperties {
 	 */
 	private Map<String, String> serverListQueryTags = new HashMap<>();
 
+	/**
+	 * Tag to query for in service list if one is not listed in serverListQueryTags.
+	 */
+	private String defaultQueryTag;
+
 	/**
 	 * Add the 'passing` parameter to /v1/health/service/serviceName.
 	 * This pushes health check passing to the server.
@@ -139,6 +144,16 @@ public class ConsulDiscoveryProperties {
 		this.hostname = this.hostInfo.getHostname();
 	}
 
+	/**
+	 *
+	 * @param serviceId The service who's filtering tag is being looked up
+	 * @return The tag the given service id should be filtered by, or null.
+     */
+	public String getQueryTagForService(String serviceId){
+		String tag = serverListQueryTags.get(serviceId);
+		return tag != null ? tag : defaultQueryTag;
+	}
+
 	public String getHostname() {
 		return this.preferIpAddress ? this.ipAddress : this.hostname;
 	}
diff --git a/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulServerList.java b/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulServerList.java
index ad70ef08..de18cb5b 100644
--- a/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulServerList.java
+++ b/spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/ConsulServerList.java
@@ -76,7 +76,7 @@ public class ConsulServerList extends AbstractServerList<ConsulServer> {
 	}
 
 	private String getTag() {
-		return this.properties.getServerListQueryTags().get(this.serviceId);
+		return this.properties.getQueryTagForService(this.serviceId);
 	}
 
 	@Override
diff --git a/spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryPropertiesTests.java b/spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryPropertiesTests.java
new file mode 100644
index 00000000..ec3ae170
--- /dev/null
+++ b/spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/discovery/ConsulDiscoveryPropertiesTests.java
@@ -0,0 +1,46 @@
+package org.springframework.cloud.consul.discovery;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.springframework.cloud.commons.util.InetUtils;
+import org.springframework.cloud.commons.util.InetUtilsProperties;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class ConsulDiscoveryPropertiesTests {
+
+    private static final String DEFAULT_TAG = "defaultTag";
+    private static final String MAP_TAG = "mapTag";
+    private static final String SERVICE_NAME_IN_MAP = "serviceNameInMap";
+    private static final String SERVICE_NAME_NOT_IN_MAP = "serviceNameNotInMap";
+    private ConsulDiscoveryProperties properties;
+    private Map<String, String> serverListQueryTags = Collections.singletonMap(SERVICE_NAME_IN_MAP, MAP_TAG);
+
+    @Before
+    public void setUp() throws Exception {
+        properties = new ConsulDiscoveryProperties(new InetUtils(new InetUtilsProperties()));
+        properties.setDefaultQueryTag(DEFAULT_TAG);
+        properties.setServerListQueryTags(serverListQueryTags);
+    }
+
+    @Test
+    public void testReturnsNullWhenNoDefaultAndNotInMap() throws Exception {
+        properties.setDefaultQueryTag(null);
+
+        assertNull(properties.getQueryTagForService(SERVICE_NAME_NOT_IN_MAP));
+    }
+
+    @Test
+    public void testGetTagReturnsDefaultWhenNotInMap() throws Exception {
+        assertEquals(DEFAULT_TAG, properties.getQueryTagForService(SERVICE_NAME_NOT_IN_MAP));
+    }
+
+    @Test
+    public void testGetTagReturnsMapValueWhenInMap() throws Exception {
+        assertEquals(MAP_TAG, properties.getQueryTagForService(SERVICE_NAME_IN_MAP));
+    }
+}
\ No newline at end of file